Tensorflow：在GPU和CPU上同时进行预测

Question

我正在使用Tensorflow，并希望通过同时使用CPU和一个GPU来加速预先训练的Keras模型的预测阶段（不关心训练阶段）。

我尝试创建两个不同的线程，分别提供两个不同的TensorFlow会话（一个在CPU上运行，另一个在GPU上运行）。每个线程以循环方式提供一定数量的批次（例如，如果总共有100个批次，则我想将20个批次分配给CPU，80个批次分配给GPU，或任何可能的组合），然后将结果组合起来。最好自动进行拆分。

然而，即使在这种情况下，似乎批次是以同步方式提供的，因为即使将少量批次发送到CPU并在GPU上计算所有其他批次（以GPU为瓶颈），我观察到总体预测时间始终高于仅使用GPU进行测试的时间。

我期望它更快，因为当只有GPU在工作时，CPU使用率约为20-30％，因此有一些CPU可用于加速计算。

我阅读了很多讨论，但它们都处理使用多个GPU进行并行处理，而不是在GPU和CPU之间进行并行处理。

以下是我编写的代码示例：在这种方式下，tensor_cpu和tensor_gpu对象从同一Keras模型中加载：

with tf.device('/gpu:0'):
    model_gpu = load_model('model1.h5')
    tensor_gpu = model_gpu(x)

with tf.device('/cpu:0'):
    model_cpu = load_model('model1.h5')
    tensor_cpu = model_cpu(x) 

然后按以下方式进行预测：

def predict_on_device(session, predict_tensor, batches):
    for batch in batches:
        session.run(predict_tensor, feed_dict={x: batch})


def split_cpu_gpu(batches, num_batches_cpu, tensor_cpu, tensor_gpu):
    session1 = tf.Session(config=tf.ConfigProto(log_device_placement=True))
    session1.run(tf.global_variables_initializer())
    session2 = tf.Session(config=tf.ConfigProto(log_device_placement=True))
    session2.run(tf.global_variables_initializer())

    coord = tf.train.Coordinator()

    t_cpu = Thread(target=predict_on_device, args=(session1, tensor_cpu, batches[:num_batches_cpu]))
    t_gpu = Thread(target=predict_on_device, args=(session2, tensor_gpu, batches[num_batches_cpu:]))

    t_cpu.start()
    t_gpu.start()

    coord.join([t_cpu, t_gpu])

    session1.close()
    session2.close()

我怎样才能实现CPU/GPU并行化？我觉得我错过了什么。

任何形式的帮助都将不胜感激！

- battuzz

---

我回答了你的问题吗？ - MWB

是的，是的，是的！！非常抱歉回复晚了，我正在忙于另一个项目，没有时间尝试这个。 我已经检查了您的代码..它可能是 intra_op_parallelism_thread 选项导致它无法工作的唯一原因吗？ - battuzz

你有什么想法可以让TensorFlow找到正确的批次数量来提供给CPU和GPU，以便我可以最小化总预测时间？ - battuzz

如果TF那样做了，我会感到惊讶。 - MWB

Answer 1

这是我的代码，演示了如何同时进行CPU和GPU执行：

这里是演示如何同时进行CPU和GPU执行的代码：

import tensorflow as tf
import numpy as np
from time import time
from threading import Thread

n = 1024 * 8

data_cpu = np.random.uniform(size=[n//16, n]).astype(np.float32)
data_gpu = np.random.uniform(size=[n    , n]).astype(np.float32)

with tf.device('/cpu:0'):
    x = tf.placeholder(name='x', dtype=tf.float32)

def get_var(name):
    return tf.get_variable(name, shape=[n, n])

def op(name):
    w = get_var(name)
    y = x
    for _ in range(8):
        y = tf.matmul(y, w)
    return y

with tf.device('/cpu:0'):
    cpu = op('w_cpu')

with tf.device('/gpu:0'):
    gpu = op('w_gpu')

def f(session, y, data):
    return session.run(y, feed_dict={x : data})


with tf.Session(config=tf.ConfigProto(log_device_placement=True, intra_op_parallelism_threads=8)) as sess:
    sess.run(tf.global_variables_initializer())

    coord = tf.train.Coordinator()

    threads = []

    # comment out 0 or 1 of the following 2 lines:
    threads += [Thread(target=f, args=(sess, cpu, data_cpu))]
    threads += [Thread(target=f, args=(sess, gpu, data_gpu))]

    t0 = time()

    for t in threads:
        t.start()

    coord.join(threads)

    t1 = time()


print t1 - t0

以下是计时结果：

• CPU线程：4-5秒（当然不同机器可能会有所不同）。

• GPU线程：5秒（它的工作量是CPU线程的16倍）。

• 同时运行两个线程：5秒

请注意，没有必要使用2个会话（但这对我也起作用）。

您可能看到不同的结果原因可能是：

• 系统资源争用（GPU执行确实会消耗一些主机系统资源，如果运行CPU线程会使其拥挤，这可能会恶化性能）

• 计时不正确

• 模型的某部分只能在GPU / CPU上运行

• 瓶颈在其他地方

• 其他问题

- MWB