我有一个在 CPU 上动态生成的数据集。样本在 python 中由一个make_sample
非常复杂的函数计算,无法转换为 tensorflow ops。因为样本生成很耗时,所以我想从多个线程调用该函数来填充输入队列。
import numpy as np
import tensorflow as tf
import time
def make_sample():
# something that takes time and needs to be on CPU w/o tf ops
p = 1
for n in range(1000000):
p = (p + np.random.random()) * np.random.random()
return np.float32(p)
read_threads = 1
with tf.device('/cpu:0'):
example_list = [tf.py_func(make_sample, [], [tf.float32]) for _ in range(read_threads)]
for ex in example_list:
ex[0].set_shape(())
batch_size = 3
capacity = 30
batch = tf.train.batch_join(example_list, batch_size=batch_size, capacity=capacity)
with tf.Session().as_default() as sess:
tf.global_variables_initializer().run()
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
try:
# dry run, left out of timing
sess.run(batch)
start_time = time.time()
for it in range(5):
print(sess.run(batch))
finally:
duration = time.time() - start_time
print('duration: {0:4.2f}s'.format(duration))
coord.request_stop()
coord.join(threads)
Run Code Online (Sandbox Code Playgroud)
令我惊讶的是,当增加 时read_threads
,CPU 使用率永远不会超过 50%。更糟糕的是,计算时间直线下降:在我的电脑上,
read_threads=1
? duration: 12s
read_threads=2
? duration: 46s
read_threads=4
? duration: 68s
read_threads=8
? duration: 112s
有没有解释,最重要的是,有一个解决方案可以在 tensorflow 上使用自定义 python 函数获得高效的多线程数据生成?
tf.py_func
重用现有的 Python 解释器。不幸的是,Python 支持并发但不支持并行。换句话说,你可以有多个 Python 线程,但任何时候只有一个可以执行 Python 代码。标准的解决方案是将生成管道移动到 TensorFlow/C++ 中,或者使用多个 Python 进程和附加层来聚合它们的结果(即使用 ZMQ 聚合来自多个 Python 进程的结果)
归档时间: |
|
查看次数: |
511 次 |
最近记录: |