python - Tensorflow 高效的每像素梯度计算

Question

我正在重新实现论文 Learning Image Matching by Simply Watching Video使用 tensorflow 时，我在从网络中获取梯度时遇到了一些严重的性能问题。为了快速回顾一下他们在论文中所做的事情，他们拥有训练有素的网络，他们做 1 个前向传播以获得插值图像，然后他们做 w*h/stride^2 反向传播以获得每个输入的输出梯度像素。由于大量的反向传播，这必须相当有效地完成，以便在合理的时间内获得梯度(在论文中，8 分钟，每个反向传播 150 毫秒乘以 128*384/16 像素(步长 4在行和列上))。由于在 tensorflow 中，由于梯度聚合，无法对多个反向传播进行批处理(例如参见 this discussion )，我需要执行以下操作:

for i in range(0, h, stride): 
    for j in range(0, w, stride):
        grad_output[0,i,j,:] = 1 #select current pixel
        grad.append(tf.gradients(predictions, images, grad_output))
        grad_output[grad_output != 0] = 0

获取每个像素的符号梯度，其中 predictions 是网络的输出张量，images 是输入，在 gpu 常量中声明为:

with tf.device('/gpu:0'):
    images = tf.constant(inp, dtype=tf.float32)

其中 inp 是包含数据的实际 numpy 数组。

单独调用 tf.gradients 需要大约 0.35 毫秒，与作者在论文中报告的时间相比，这已经太多了。但是最多的时间花在评估符号梯度上，比如:

for i in range(0, len(grad)):
    res = sess.run(grad[i])

这大约需要 1.5 秒，非常慢。现在，随后调用 sess.run(grad[i])(具有相同的索引 i)非常快，大约 100 毫秒，同时运行 for 循环改变 i 在每次迭代中导致每次迭代大约 1.5 秒。看到这种行为后，我猜测将东西移动到 GPU 的开销很大，这可能吗？如果是这种情况，我该如何避免呢？我已经将 images 张量移动到 GPU 常量，而不是使用占位符并依赖 sess.run 中的 feed_dict，但这并没有'对性能有任何明显的影响。有什么想法可以加快对符号梯度的评估吗？我觉得我在这里遗漏了一些简单的东西，因为 1 个反向传播需要 1.5 秒，这与任何现实场景都相去甚远(例如，训练网络能够每秒处理大约 100 个样本，所以我猜这不是架构问题..)

谢谢!

Answer 1

这是我的想法:

1. 尝试可视化学习图，例如通过tensorboard和 SummaryWriter .循环中的 tf.gradients 看起来很可疑——确保您创建给定张量的次数不要超过必要的次数。

2. 如果它没有改变您的语义，请尝试一次计算所有梯度:

   res = sess.run(grad)

假设 grad 是一个张量列表。在循环中执行 sess.run 将多次重新计算 grad[i] 和 grad[j] 的任何公共(public)父级。

希望对您有所帮助!