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)
...
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成功地完成了这个.我还认为这将是对优化器已经构建的计算图的略微修改.

如何更快地完成,或者是否有任何其他方法来计算损失的梯度w.r.t.投入？

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