我正在使用带有张量流的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 Build /编译/训练我的模型 .
我曾经想过,也许我只需要使用python的多处理模块,并为每个运行 predict_proba(batch_n)
的gpu启动一个进程 . 我知道这在理论上是可能的,因为我的另一个SO帖子:Keras + Tensorflow and Multiprocessing in Python . 然而,这仍然让我陷入两难境地:不知道如何实际操作这个过程 .
我的问题归结为:当使用Tensorflow作为Keras的后端时,如何将Keras中的一个模型的预测与多个gpus并行化?
另外,我很好奇是否只有一个gpu可以进行类似的预测并行化 .
高级描述或代码示例将不胜感激!
谢谢!
2 回答
我创建了一个简单的例子来展示如何跨多个gpus运行keras模型 . 基本上,创建了多个进程,每个进程都拥有一个gpu . 要指定进程中的gpu id,设置env变量CUDA_VISIBLE_DEVICES是一种非常简单的方法(os.environ [“CUDA_VISIBLE_DEVICES”]) . 希望这个git repo可以帮到你 .
https://github.com/yuanyuanli85/Keras-Multiple-Process-Prediction
您可以使用此功能来并行化Keras模型(信用到kuza55) .
https://github.com/kuza55/keras-extras/blob/master/utils/multi_gpu.py
.