Keras + Tensorflow：多 GPU 上的预测

Question

我正在使用带有 tensorflow 作为后端的 Keras。我有一个 compiled/trained 型号。

我的预测循环很慢，所以我想找到一种方法来并行化 predict_proba 调用以加快速度。我想获取（数据）批次的列表，然后每个可用的 gpu，运行 model.predict_proba() 在这些批次的一个子集上。
本质上：

data = [ batch_0, batch_1, ... , batch_N ]
on gpu_0 => return predict_proba(batch_0)
on gpu_1 => return predict_proba(batch_1)
...
on gpu_N => return predict_proba(batch_N)

我知道在纯 Tensorflow 中可以将操作分配给给定的 gpu (https://www.tensorflow.org/tutorials/using_gpu)。但是，我不知道这如何转化为我的情况，因为我已经使用 Keras 的 api.

构建/compiled/trained 我的模型

我曾想也许我只需要使用 python 的多处理模块并为每个 gpu 启动一个进程运行 predict_proba(batch_n)。我知道这在理论上是可能的，因为我的另一个 SO post：。然而，这仍然让我陷入不知道如何实际 "choose" 一个 gpu 来操作进程的困境。

我的问题归结为：当使用 Tensorflow 作为 Keras 的后端时，如何在 Keras 中跨多个 GPU 并行预测一个模型？

此外，我很好奇是否仅用一个 gpu 就可以进行类似的预测并行化。

非常感谢高级描述或代码示例！

谢谢！

Answer 1

我创建了一个简单示例来展示如何运行跨多个 GPU 的 keras 模型。基本上，创建了多个进程，每个进程都拥有一个 gpu。要在进程中指定 gpu id，设置 env 变量 CUDA_VISIBLE_DEVICES 是一种非常简单的方法 (os.environ["CUDA_VISIBLE_DEVICES"])。希望这个 git 回购可以帮助你。

https://github.com/yuanyuanli85/Keras-Multiple-Process-Prediction

Answer 2

您可以使用此函数并行化 Keras 模型（归功于 kuza55）。
https://github.com/kuza55/keras-extras/blob/master/utils/multi_gpu.py
.

from keras.layers import merge
from keras.layers.core import Lambda
from keras.models import Model

import tensorflow as tf

def make_parallel(model, gpu_count):
    def get_slice(data, idx, parts):
        shape = tf.shape(data)
        size = tf.concat([ shape[:1] // parts, shape[1:] ],axis=0)
        stride = tf.concat([ shape[:1] // parts, shape[1:]*0 ],axis=0)
        start = stride * idx
        return tf.slice(data, start, size)

    outputs_all = []
    for i in range(len(model.outputs)):
        outputs_all.append([])

    #Place a copy of the model on each GPU, each getting a slice of the batch
    for i in range(gpu_count):
        with tf.device('/gpu:%d' % i):
            with tf.name_scope('tower_%d' % i) as scope:

                inputs = []
                #Slice each input into a piece for processing on this GPU
                for x in model.inputs:
                    input_shape = tuple(x.get_shape().as_list())[1:]
                    slice_n = Lambda(get_slice, output_shape=input_shape, arguments={'idx':i,'parts':gpu_count})(x)
                    inputs.append(slice_n)                

                outputs = model(inputs)

                if not isinstance(outputs, list):
                    outputs = [outputs]

                #Save all the outputs for merging back together later
                for l in range(len(outputs)):
                    outputs_all[l].append(outputs[l])

    # merge outputs on CPU
    with tf.device('/cpu:0'):
        merged = []
        for outputs in outputs_all:
            merged.append(merge(outputs, mode='concat', concat_axis=0))

        return Model(input=model.inputs, output=merged)

Keras + Tensorflow：多 GPU 上的预测

Keras + Tensorflow: Prediction on multiple gpus

python

python-multiprocessing

keras

tensorflow