python - 是否可以在 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]), ...

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

使用 TensorFlow 的自动微分功能，我只需调用 f_prime.eval(feed_dict={x: xi}) 即可计算 f'(x1)、f'(x2)、...、f'(xn)对于 (x1, x2, ..., xn) 中的每个值 xi 一次。这根本没有效率:我想使用矢量化形式，但我不知道如何做到这一点。

也许使用tf.stop_gradient()不知何故？或使用 grad_ys tf.gradients() 中的参数？

Answer 1

经过一番挖掘，在 TensorFlow 中计算每个示例的梯度似乎并不简单，因为该库执行标准反向传播来计算梯度(其他深度学习库如 PyTorch、Theano 等也是如此) on)，它从未实际计算每个示例的梯度，它直接获得每个示例的梯度之和。查看this discussion了解更多详情。

但是，有一些技术可以解决这个问题，至少对于某些用例来说是这样。例如，论文Efficient per-example gradient computation Ian Goodfellow 解释了如何有效计算包含平方导数之和的每个示例向量。以下是显示计算的论文摘录(但我强烈建议您阅读这篇论文，它很短):

该算法的复杂度为 O(mnp) 而不是 O(mnp²)，其中 m 是示例数，n 是神经网络中的层数，p 是每层神经元的数量。因此，它比简单的方法(即每个示例执行一次反向传播)要快得多，尤其是当 p 很大时，甚至在使用 GPU 时速度更快(这会大大加快矢量化方法的速度)。