Python Keras+；Tensorflow：多GPU上的预测

我使用Keras和tensorflow作为后端。 我有一个经过编译/训练的模型

我的预测循环很慢，所以我想找到一种方法来并行化

predict\u 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中，可以将ops分配给给定的gpu（）。然而，我不知道这如何解释我的情况，因为我已经使用Keras的api构建/编译/训练了我的模型

我曾想，也许我只需要使用python的多处理模块，并根据gpu启动一个进程，该进程将运行

predict\u proba（batch\u n）

。我知道这在理论上是可能的，因为我的另一篇文章是：。然而，这仍然让我左右为难，不知道如何真正“选择”一个gpu来操作这个过程

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

此外，我很好奇，如果类似的并行预测是可能的，只有一个gpu

如能提供高级描述或代码示例，将不胜感激

---

谢谢

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

---

您可以使用此功能并行化Keras模型（贷记到）。

---

经过训练并将模型保存为json，当从json和prediction重新加载模型时，它失败了，因为不知道tf，如何修复它，或者最好的解决方案是什么？是的，这是代码中已知的问题。当您运行此函数时，您将更改图形。因此，您只需要保存一个分支，就可以在以后加载它，以便在单个gpu上进行预测

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)