使用 TF-Slim 的全卷积 ResNet 运行速度非常慢

Question

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

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，但运行时仍然大致相同。

更新2。看起来我想出了速度的问题。我进行全卷积训练，我的图像具有任意大小和纵横比。如果我提供固定大小的虚拟随机 numpy 张量，它会运行得很快。如果生成 10 个预定义大小的输入张量，前 10 次迭代很慢，但随后会加快。看起来在 TensorFlow 中，每次迭代时调整所有张量的大小不如在 Caffe 中有效。我将在项目的 GitHub 上提交工单。

Answer 1

问题是由于输入图像的任意大小造成的。 TensorFlow 有一种称为自动调整的功能，因此在运行时它们会针对每个特定的输入大小分析各种算法，并决定哪种算法最好。就我而言，这发生在每次迭代中。

解决方法是设置环境变量TF_CUDNN_USE_AUTOTUNE=0：

export TF_CUDNN_USE_AUTOTUNE=0
python myscript.py

此 Github 门票的更多信息：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, l 中完成的。 172. 它们被直接添加到run 操作的“获取”中，因此并行和尽快执行。