Keras + Tensorflow：多GPU预测

Question

我正在使用Keras和TensorFlow作为后端。 我有一个已编译/训练好的模型。

我的预测循环速度较慢，因此我希望找到一种方法来并行化`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构建/编译/训练模型的，所以我不知道如何转换到我的情况。

我曾认为，也许我只需要使用Python的多进程模块，并启动每个GPU的一个进程来运行predict_proba(batch_n)。我知道这在理论上是可能的，参考我的另一个SO帖子：Keras + Tensorflow和Python中的多处理。但是，这仍然让我陷入了一个困境，即不知道如何实际上“选择”一个GPU来操作该进程。

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

此外，我很想知道是否可以在仅使用一个GPU的情况下进行类似的预测并行化。

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

谢谢！

- John Cast

Answer 1

我创建了一个简单的示例，展示如何在多个GPU上运行keras模型。基本上，会创建多个进程，每个进程拥有一个GPU。要在进程中指定GPU ID，设置环境变量CUDA_VISIBLE_DEVICES是一种非常直接的方法（os.environ["CUDA_VISIBLE_DEVICES"]）。希望这个git仓库可以帮到你。 https://github.com/yuanyuanli85/Keras-Multiple-Process-Prediction

- VictorLi

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)

- Temak

---

训练并将模型保存为json后，重新从json加载模型进行预测时，出现了tf未知的错误。如何解决或找到最佳解决方案？ - bygreencn

是的，这是此代码已知的问题。当您运行此函数时，会更改图形。因此，您需要仅保存一个分支，以便在单个GPU上进行预测时能够加载它。 - Temak