Keras + Tensorflow:预测多个gpus

Question

我正在使用带有张量流的Keras作为后端.我有一个编译/训练的模型.

我的预测循环很慢,所以我想找到一种方法来并行化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可以将ops分配给给定的gpu(https://www.tensorflow.org/tutorials/using_gpu).但是,我不知道这是如何转化为我的情况,因为我使用Keras的api构建/编译/训练了我的模型.

我曾经想过,也许我只需要使用python的多处理模块并开始运行每个gpu的进程predict_proba(batch_n).我知道这在理论上是可能的,因为我的另一个SO帖子:Keras + Tensorflow和Python中的多处理.然而,这仍然让我不知道如何实际"选择"一个gpu来操作这个过程.

我的问题归结为:当使用Tensorflow作为Keras的后端时,如何将Keras中的一个模型的预测与多个gpus并行化？

另外,我很好奇是否只有一个gpu可以进行类似的预测并行化.

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

谢谢!

Answer 1

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

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