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.

如何更快地制作,或者是否有其他方法来计算输入损失的梯度？

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