python-3.x - 在 Tensorflow 中计算雅可比和导数非常慢

Question

是否有一种更有效的方法来计算雅可比行列式（必须有，它甚至不会针对单个批次运行）我想计算自解释神经网络中给出的损失。输入的形状为 (32, 365, 3)，其中 32 是批量大小。我想最小化的损失是论文的公式 3 。

我相信我没有以最佳方式使用 GradientTape。

def compute_loss_theta(tape, parameter, concept, output, x):

    b = x.shape[0]
    in_dim = (x.shape[1], x.shape[2])

    feature_dim = in_dim[0]*in_dim[1]

    J = tape.batch_jacobian(concept, x)
    
    grad_fx = tape.gradient(output, x)
    grad_fx = tf.reshape(grad_fx,shape=(b, feature_dim))
    J = tf.reshape(J, shape=(b, feature_dim, feature_dim))

    parameter = tf.expand_dims(parameter, axis =1)

    loss_theta_matrix = grad_fx - tf.matmul(parameter, J)

    loss_theta = tf.norm(loss_theta_matrix)

    return loss_theta


for i in range(10):
    for x, y in train_dataset:

        with tf.GradientTape(persistent=True) as tape:
            tape.watch(x)
            
            parameter, concept, output = model(x)

            loss_theta = compute_loss_theta(tape, parameter, concept, output , x)

            loss_y = loss_object(y_true=y, y_pred=output)
                
            loss_value = loss_y + eps*loss_theta

        gradients = tape.gradient(loss_value, model.trainable_weights)
        optimizer.apply_gradients(zip(gradients, model.trainable_weights))

标签： python-3.xoptimizationtensorflow2.0