在 TensorFlow 中批量访问单个渐变的最佳方法是什么？

Question

我目前正在分析梯度如何在使用 Tensorflow 2.x 训练 CNN 的过程中发展。我想要做的是将批次中的每个梯度与整个批次的梯度进行比较。目前我对每个训练步骤都使用这个简单的代码 sn-p：

[...]
loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
[...]

# One training step
# x_train is a batch of input data, y_train the corresponding labels
def train_step(model, optimizer, x_train, y_train):

    # Process batch
    with tf.GradientTape() as tape:
        batch_predictions = model(x_train, training=True)
        batch_loss = loss_object(y_train, batch_predictions)
    batch_grads = tape.gradient(batch_loss, model.trainable_variables)
    # Do something with gradient of whole batch
    # ...

    # Process each data point in the current batch
    for index in range(len(x_train)):
        with tf.GradientTape() as single_tape:
            single_prediction = model(x_train[index:index+1], training=True)
            single_loss = loss_object(y_train[index:index+1], single_prediction)
        single_grad = single_tape.gradient(single_loss, model.trainable_variables)
        # Do something with gradient of single data input
        # ...

    # Use batch gradient to update network weights
    optimizer.apply_gradients(zip(batch_grads, model.trainable_variables))

    train_loss(batch_loss)
    train_accuracy(y_train, batch_predictions)

我的主要问题是单手计算每个梯度时计算时间会爆炸，尽管在计算批次梯度时这些计算应该已经由 Tensorflow 完成。原因是GradientTape 和compute_gradients 总是返回一个梯度，无论给出单个还是多个数据点。所以这个计算必须针对每个数据点进行。

我知道我可以通过使用为每个数据点计算的所有单个梯度来计算批次的梯度来更新网络，但这在节省计算时间方面只起次要作用。

有没有更有效的方法来计算单个梯度？

Answer 1

您可以使用梯度带的jacobian 方法获得雅可比矩阵，这将为您提供每个单独损失值的梯度：

import tensorflow as tf

# Make a random linear problem
tf.random.set_seed(0)
# Random input batch of ten four-vector examples
x = tf.random.uniform((10, 4))
# Random weights
w = tf.random.uniform((4, 2))
# Random batch label
y = tf.random.uniform((10, 2))
with tf.GradientTape() as tape:
    tape.watch(w)
    # Prediction
    p = x @ w
    # Loss
    loss = tf.losses.mean_squared_error(y, p)
# Compute Jacobian
j = tape.jacobian(loss, w)
# The Jacobian gives you the gradient for each loss value
print(j.shape)
# (10, 4, 2)
# Gradient of the loss wrt the weights for the first example
tf.print(j[0])
# [[0.145728424 0.0756840706]
#  [0.103099883 0.0535449386]
#  [0.267220169 0.138780832]
#  [0.280130595 0.145485848]]