是否可以在TensorFlow中有效地计算每个示例的梯度，只需一次图表运行？

Question

TD; DR：有没有办法以矢量化的形式在一次图表运行中评估f'（x1），f'（x2），...，f'（xn）？其中f'（x）是f（x）的导数。

类似的东西：

x = tf.placeholder(tf.float32, shape=[100])
f = tf.square(x)
f_grad = tf.multiple_gradients(x) # f_grad contains f'(x[0]), f'(x[1]), ...

更具体地说，我正在尝试手动实现Black Box随机变分推理（BBSVI）（我知道我可以使用像Edward这样的库，但我正在尝试自己实现它）。 有一点，我需要计算跨越x（x1，x2，...，xn）的许多不同值的f'（x）g（x）的平均值，其中f（x）和g（x）是两个函数，f'（x）是f（x）的导数。

使用TensorFlow的autodiff功能，我可以通过简单地为{x1}中的每个值xi调用f_prime.eval(feed_dict={x: xi})来计算f'（x1），f'（x2），...，f'（xn）， x2，...，xn）。这根本不高效：我想使用矢量化形式，但我不知道该怎么做。

也许以某种方式使用tf.stop_gradient()？或者使用grad_ys中的tf.gradients()参数？

Answer 1

经过一些挖掘后，似乎在TensorFlow中计算每个示例的渐变并不容易，因为这个库执行标准的反向传播来计算渐变（就像PyTorch，Theano等其他深度学习库一样） on），它从不实际计算每个例子的梯度，它直接获得每个例子梯度的总和。查看this discussion了解详情。

但是，有一些技术可以解决这个问题，至少在某些用例中是这样。例如，Ian Goodfellow的论文Efficient per-example gradient computation解释了如何有效地计算包含平方导数之和的每个例子向量。以下是显示计算的论文的摘录（但我强烈建议您阅读论文，它非常简短）：

该算法是O（mnp）而不是O（mnp²），其中m是例子的数量，n是神经网络中的层数，p是每层神经元的数量。所以它比天真的方法快得多（即，每个例子执行一次反向支持），特别是当p很大时，甚至更多的时候使用GPU（这会大大加速矢量化方法）。

Answer 2

您可以使用tf.vectorized_map(forward_and_backward_fn, batch_of_inputs)有效地计算每个示例的梯度。

Answer 3

在Tensorflow上，它准确地演示了每个示例渐变的示例，如下所示：

# Computing per-example gradients
batch_size = 10
num_features = 32
layer = tf.keras.layers.Dense(1)

def model_fn(arg):
  with tf.GradientTape() as g:
    inp, label = arg
    inp = tf.expand_dims(inp, 0)
    label = tf.expand_dims(label, 0)
    prediction = layer(inp)
    loss = tf.nn.l2_loss(label - prediction)
  return g.gradient(loss, (layer.kernel, layer.bias))

inputs = tf.random.uniform([batch_size, num_features])
labels = tf.random.uniform([batch_size, 1])
per_example_gradients = tf.vectorized_map(model_fn, (inputs, labels))
assert per_example_gradients[0].shape == (batch_size, num_features, 1)
assert per_example_gradients[1].shape == (batch_size, 1)

您可以使用vectorized_map参考the official link。