python - 凯拉斯 + tensorflow : Prediction on multiple gpus

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 and Multiprocessing in Python .然而，这仍然让我陷入不知道如何实际“选择”一个 gpu 来运行该进程的困境。

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

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

将不胜感激高级描述或代码示例!

谢谢!

Answer 1

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

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