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

Question

我正在将最初在 Caffe 中实现的像素标记(FCN 样式)代码移植到 TensorFlow。我使用步幅为 16px 的 Slim ResNets (ResNet-101) 实现，并使用上卷积层进一步对其进行上采样，以实现最终步幅为 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