TensorFlow：在输入中获取梯度时性能低下

Question

我正在使用 TensorFlow 构建一个简单的多层感知器，我还需要获取神经网络输入端损失的梯度（或误差信号）。

这是我的有效代码：

cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(self.network, self.y))
optimizer = tf.train.AdagradOptimizer(learning_rate=nn_learning_rate).minimize(cost)
...
for i in range(epochs):
    ....
    for batch in batches:
        ...
        sess.run(optimizer, feed_dict=feed_dict)
        grads_wrt_input = sess.run(tf.gradients(cost, self.x), feed_dict=feed_dict)[0]

（经过编辑以包含训练循环）

没有最后一行 (grads_wrt_input...)，它在 CUDA 机器上运行得非常快。但是，tf.gradients() 将性能大大降低十倍或更多。

我记得在反向传播算法中，节点处的误差信号被计算为中间值，我已经使用 Java 库 DeepLearning4j 成功地完成了这项工作。我还觉得这将是对 optimizer.

已经构建的计算图的轻微修改

如何让它更快，或者是否有任何其他方法来计算损失的梯度 w.r.t。输入？

Answer 1

tf.gradients() 函数在每次调用时都会构建一个新的反向传播图，因此速度变慢的原因是 TensorFlow 必须在循环的每次迭代中解析一个新图。（这可能非常昂贵：当前版本的 TensorFlow 针对多次执行相同图进行了优化。）

幸运的是，解决方案很简单：只需在循环外计算一次梯度。您可以按如下方式重组代码：

cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(self.network, self.y))
optimizer = tf.train.AdagradOptimizer(learning_rate=nn_learning_rate).minimize(cost)
grads_wrt_input_tensor = tf.gradients(cost, self.x)[0]
# ...
for i in range(epochs):
    # ...
    for batch in batches:
        # ...
        _, grads_wrt_input = sess.run([optimizer, grads_wrt_input_tensor],
                                      feed_dict=feed_dict)

请注意，为了提高性能，我还合并了两个 sess.run() 调用。这确保了前向传播和大部分反向传播将被重用。

顺便说一句，查找此类性能错误的一个技巧是在开始训练循环之前调用 tf.get_default_graph().finalize()。如果您无意中将任何节点添加到图中，这将引发异常，从而更容易追踪这些错误的原因。

TensorFlow：在输入中获取梯度时性能低下

TensorFlow: slow performance when getting gradients at inputs

tensorflow