Tensorflow多个样本,无需批处理

时间:2017-02-28 04:06:37

标签: python multithreading tensorflow

我有一个图表,不幸的是有一些节点不支持批处理。(自定义操作尚未充实)我已经能够有多个线程调用sess.run并通过feed_dict抛出数据。我现在已经将我的数据转换为tfrecords以正确使用队列,但仍然无法找到一种方法来告诉它并行运行多个图形实例而不需要多个线程调用sess.run()。我假设tensorflow开发人员已经在某个地方创建了一种更“pythonic”的方式,但我还没有找到它。我如何在tensorflow中执行此操作?

编辑:即使批量处理数据,前一个问题仍然存在,因为我的计算花费了一半的时间在cpu上,而一半用在gpu上,因此不管批处理,一个人将等待另一半的时间。我想让图表异步训练多个样本以填充该空间。

编辑2:我想我必须在这里为那些不想阅读上述文字的人提供伪代码。

import tensorflow as tf

resultOfCPUCalculation = someCPUOnlyOP(inputData)\\does not support batching
gpuResults = aBunchOfGPUOps(resultOfCPUCalculation)
with tf.Session() as sess:
    sess.run([gpuResults])
    //only uses 1 cpu core, and the gpu is idle while it's doing it's thing.

我想以“管道”方式执行此操作,只要cpu操作完成,它就会从另一个样本开始。

1 个答案:

答案 0 :(得分:0)

import tensorflow as tf

input_example = get_input_example()
cpu_output = some_cpu_only_op(input_example)

cpu_output_batch = tf.train.batch(input_example, batch_size, num_threads)
gpu_output_batch = a_bunch_of_gpu_ops(cpu_output_batch)

with tf.Session() as sess:
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(coord=coord)

    for i in range(num_train_steps):
        output_values = sess.run(gpu_output_batch)
        do_stuff_with(output_values)

    coord.request_stop()
    coord.join(threads)

input_example必须由队列提供。