我试图建立一个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进行批处理,这使事情变得更糟。我试着在运行循环之前添加一个小睡眠,只是为了确保队列已经填满 - 但没有运气。
非常感谢任何帮助。