Tensorflow:输入管道非常慢/不缩放

时间:2017-04-09 15:36:59

标签: multithreading tensorflow batch-processing

我试图建立一个Tensorflow输入管道,用于将图像输入AlexNet以进行特征提取(不是用于训练,这是一个问题)。由于AlexNet相当小,因此必须以高速率提供输入数据以实现可接受的性能(约1000张图像/秒)。

我的图像是400x300 JPEG,平均每张图像24KB。

不幸的是,似乎Tensorflow输入管道无法跟上运行AlexNet的GTX1080。

我的输入管道很简单:加载文件,解码图像,调整大小并批量处理。

我创建了一个小基准来显示问题:

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import time
import glob
import os

IMAGE_DIR = 'images'
EPOCHS = 1


def main():
    print('batch_size\tnum_threads\tms/image')

    for batch_size in [16, 32, 64, 128]:
        for num_threads in [1, 2, 4, 8]:
            run(batch_size, num_threads)


def run(batch_size, num_threads):
    filenames = glob.glob(os.path.join(IMAGE_DIR, '*.jpg'))
    (filename,) = tf.train.slice_input_producer(
        [filenames],
        capacity=2 * batch_size * num_threads,
        num_epochs=EPOCHS)

    raw = tf.read_file(filename)
    decoded = tf.image.decode_jpeg(raw, channels=3)
    resized = tf.image.resize_images(decoded, [227, 227])

    batch = tf.train.batch(
        [resized],
        batch_size,
        num_threads,
        2 * batch_size * num_threads,
        enqueue_many=True)

    init_op = tf.group(
        tf.global_variables_initializer(),
        tf.local_variables_initializer())

    with tf.Session() as sess:
        sess.run(init_op)

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        t = time.time()
        try:
            while not coord.should_stop():
                sess.run(batch)
        except tf.errors.OutOfRangeError:
            pass
        finally:
            coord.request_stop()

        tpe = (time.time() - t) / (len(filenames) * EPOCHS) * 1000
        print('{: <11}\t{: <10}\t{: <8}'
              .format(batch_size, num_threads, tpe))

        coord.join(threads)


if __name__ == "__main__":
    main()

在MacBook Pro上运行(2015年初,2,9 GHz Intel Core i5)会产生以下结果:

batch_size      num_threads     ms/image
16              1               4.81571793556
16              2               3.00584602356
16              4               2.94281005859
16              8               2.94555711746
32              1               3.51123785973
32              2               1.82255005836
32              4               1.85884213448
32              8               1.88741898537
64              1               2.9537730217
64              2               1.58108997345
64              4               1.57125210762
64              8               1.57615303993
128             1               2.71797513962
128             2               1.67120599747
128             4               1.6521999836
128             8               1.6885869503

它显示整体不良性能,远远低于每张图像1 / ms。此外,它不会超出两个线程,在这种情况下是预期的,因为它只是一个双核处理器。

在具有24个内核的2.5Ghz AMD Opteron 6180 SE上运行相同的基准测试产生以下结果:

batch_size      num_threads     ms/image
16              1               13.983194828
16              2               6.80965399742
16              4               6.67097783089
16              8               6.63090395927
32              1               12.0395629406
32              2               5.72535085678
32              4               4.94155502319
32              8               4.99696803093
64              1               10.9073989391
64              2               4.96317911148
64              4               3.76832485199
64              8               3.82816386223
128             1               10.2617599964
128             2               5.20488095284
128             4               3.16122984886
128             8               3.51550602913

在这里,单线程/整体性能也非常差,并且它不会超过2/4线程。

在任何情况下,系统既不是IO也不是CPU。对于两个系统,使用OpenCV加载和调整图像大小可以提供更好的数字(MacBook中的~0.86ms /图像,在这种情况下是CPU绑定的,服务器上的图像最多约为0.22ms /图像,在这种情况下是IO绑定)。

Tensorflow在这里发生了什么?我怎样才能加快速度呢?

我已经尝试手动组装一批图像并使用enqueue_many进行批处理,这使事情变得更糟。我试着在运行循环之前添加一个小睡眠,只是为了确保队列已经填满 - 但没有运气。

非常感谢任何帮助。

0 个答案:

没有答案