Tensorflow高效的逐像素梯度计算

Question

我正在使用tensorflow重新实现通过观看视频学习图像匹配这篇论文，但是在从网络中获取梯度时遇到了一些严重的性能问题。简要回顾一下他们在论文中所做的，他们有训练好的网络，他们进行一次前向传递以获取插值图像，然后他们进行w*h/stride^2次反向传播，以获取每个像素点输出相对于输入的梯度。由于反向传播的数量很高，因此必须以相当高的效率来完成，以便在合理的时间内获得梯度（在论文中，8分钟，每个反向传播150毫秒，128*384/16像素（行和列上都是步长4））。由于在tensorflow中无法批量处理多个反向传播，因为需要梯度聚合（例如此讨论），因此我需要做类似于以下的操作：

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上的常量a：

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秒，非常慢。现在，对于相同的索引 i ，后续调用sess.run(grad [i])非常快，大约100毫秒，而在每次迭代中更改i的for循环的运行结果为每次迭代大约1.5秒。看到这种行为后，我的猜测是在将东西移动到GPU上存在很大的开销，这可能吗？如果是这种情况，我该如何避免它？我已经将images张量移动到GPU常量中，而不是使用占位符并依赖feed_dict在sess.run中，但这对性能没有任何明显影响。有什么方法可以加速符号梯度的评估吗？我觉得我错过了一些简单的东西，因为1个反向传播需要1.5秒，这真的离任何现实场景都很远（例如，训练网络能够处理大约100个样本每秒，所以我想这不是一个架构问题..）谢谢！

- powder

Answer 1

以下是我的想法：

1. 尝试通过tensorboard和SummaryWriter来可视化学习图表。在循环中使用tf.gradients看起来有些可疑，确保您不会比必要的次数更多地创建给定的张量。

2. 如果不改变语义，请尝试一次计算所有梯度：

   res = sess.run(grad)

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

希望这可以帮助到您！

- sygi