使用TF-Slim的全卷积ResNets运行非常慢

Question

我正在将最初在Caffe中实现的像素标记(FCN-样式)代码移植到TensorFlow上。我使用Slim实现的ResNets (ResNet-101)的步长为16 8px，并进一步对其进行向上卷积层的采样，以达到8 8px的最终步长。我有batch_size=1，因为输入的图像是任意大小的。问题是训练真的很慢。它在3.5分钟内处理了100幅图像，而我最初的caffe实现在相同的硬件上用30秒完成了(特斯拉K40m)。下面是我所拥有的代码的简化版本：

import datetime as dt

import tensorflow as tf
import tensorflow.contrib.slim as slim
from tensorflow.contrib.slim.nets import resnet_v1

from MyDataset import MyDataset
from TrainParams import TrainParams

dataset = MyDataset()
train_param = TrainParams()

#tf.device('/gpu:0')

num_classes = 15

inputs = tf.placeholder(tf.float32, shape=[1, None, None, 3])

with slim.arg_scope(resnet_v1.resnet_arg_scope(False)):
    mean = tf.constant([123.68, 116.779, 103.939],
                       dtype=tf.float32, shape=[1, 1, 1, 3], name='img_mean')
    im_centered = inputs - mean
    net, end_points = resnet_v1.resnet_v1_101(im_centered,
                                              global_pool=False, output_stride=16)

    pred_upconv = slim.conv2d_transpose(net, num_classes,
                                        kernel_size = [3, 3],
                                        stride = 2,
                                        padding='SAME')

    targets = tf.placeholder(tf.float32, shape=[1, None, None, num_classes])

    loss = slim.losses.sigmoid_cross_entropy(pred_upconv, targets)


log_dir = 'logs/'

variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"])
restorer = tf.train.Saver(variables_to_restore)

with tf.Session() as sess:

  sess.run(tf.initialize_all_variables())
  sess.run(tf.initialize_local_variables())

  restorer.restore(sess, '/path/to/ResNet-101.ckpt')

  optimizer = tf.train.GradientDescentOptimizer(learning_rate=.001)
  train_step = optimizer.minimize(loss)
  t1 = dt.datetime.now()
  for it in range(10000):
      n1=dt.datetime.now()
      batch = dataset.next_batch() # my function that prepares training batch
      sess.run(train_step, feed_dict={inputs: batch['inputs'],
                                      targets: batch['targets']})
      n2=dt.datetime.now()
      time = (n2-n1).microseconds/(1000)
      print("iteration ", it, "time", time)

我只是在学习这个框架，我只在几天内把这些代码组合在一起，所以我知道它可能不是最好的。正如您所看到的，我还试图测量数据准备代码和网络前向后传递所需的实际时间。与实际运行时相比，这一次实际上要小得多，加起来100个迭代，只有50秒。我怀疑可能有一些线程/进程同步正在进行，这是无法测量的，但我发现这很奇怪。top命令显示了大约10个进程，标题与它可能生成的主进程相同。我也收到这样的警告：

I tensorflow/core/common_runtime/gpu/pool_allocator.cc:245] PoolAllocator: After 1692 get requests, put_count=1316 evicted_count=1000 eviction_rate=0.759878 and unsatisfied allocation rate=0.87234
I tensorflow/core/common_runtime/gpu/pool_allocator.cc:257] Raising pool_size_limit_ from 100 to 110

你能不能给我指出一个正确的方向，让我加快速度？

谢谢。

最新消息。经过更多的研究后，我发现与队列相比，“喂食”数据的速度可能比较慢，所以我在一个单独的线程：https://gist.github.com/eldar/0ecc058670be340b92e5a1044dc8a089中重新实现了队列代码：，但是运行时仍然是相同的。

UPDATE2。看来我知道速度有什么问题了。我全身心地训练，我的图像大小和高宽比都是任意的.如果我给出的是固定尺寸的随机无规张量，它就会工作得很快。如果生成10个预定义大小的输入张量，那么前10个迭代是缓慢的，但是它会加速。在TensorFlow中，每次迭代时对所有张量的调整都不如在Caffe中有效。我将在该项目的GitHub上提交一张票证。

Answer 1

这一问题是由于输入任意大小的图像造成的。TensorFlow有一个叫做自动调优的东西，所以在运行时，他们会为每个特定的输入大小分析各种算法，并决定哪一个是最好的。在我的例子中，每次迭代都会发生这种情况。

解决方案是设置环境变量TF_CUDNN_USE_AUTOTUNE=0。

export TF_CUDNN_USE_AUTOTUNE=0
python myscript.py

更多信息请看这张吉特布票：https://github.com/tensorflow/tensorflow/issues/5048。

Answer 2

一般来说，TensorFlow resnet的实现不应该比caffe慢(很多)。我刚才比较了caffe/barrista (https://github.com/classner/barrista/tree/master/examples/residual-nets)中的实现和Tensorflow (https://github.com/tensorflow/models/tree/master/resnet)中的例子，它们在相同速度下的完全训练中的差异可以忽略不计。

首先，我在Tensorflow实现方面遇到了麻烦，这让我来到了这个页面。原因是，我构建的github版本不稳定，因为中间开发代码非常慢。一个git pull和重新编译解决了这个问题。

但是，如果您正在为自己重新实现，请注意如何触发BatchNorm更新操作。在tensorflow示例中，这是在resnet_model.py，1.172中完成的。它们直接添加到run操作的“获取”中，因此并行和快速地执行。