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 成功地完成了此操作。我还认为这将是对优化器已经构建的计算图的轻微修改。

如何才能使其更快，或者是否有其他方法来计算损失的梯度。输入？

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()在开始训练循环之前。如果您无意中向图表添加任何节点，这将引发异常，从而更容易追踪这些错误的原因。