Keras + Tensorflow:对多 GPU 的预测 [英] Keras + Tensorflow: Prediction on multiple gpus

问题描述

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

I'm using Keras with tensorflow as backend. I have one compiled/trained model.

我的预测循环很慢，所以我想找到一种方法来并行化 predict_proba 调用以加快速度.我想获取一个批次(数据)列表，然后根据可用的 GPU，在这些批次的一个子集上运行 model.predict_proba().
本质上:

My prediction loop is slow so I would like to find a way to parallelize the predict_proba calls to speed things up. I would like to take a list of batches (of data) and then per available gpu, run model.predict_proba() over a subset of those batches.
Essentially:

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 构建/编译/训练了我的模型.

I know that it's possible in pure Tensorflow to assign ops to a given gpu (https://www.tensorflow.org/tutorials/using_gpu). However, I don't know how this translates to my situation since I've built/compiled/trained my model using Keras' api.

我曾想，也许我只需要使用 python 的多处理模块并为每个 gpu 启动一个进程，该进程将运行 predict_proba(batch_n).我知道这在理论上是可能的，因为我的另一个 SO 帖子是:Keras + Tensorflow 和 Python 中的多处理.然而，这仍然让我陷入两难境地，不知道如何真正选择"一个 GPU 来运行这个过程.

I had thought that maybe I just needed to use python's multiprocessing module and start a process per gpu that would run predict_proba(batch_n). I know this is theoretically possible given another SO post of mine: Keras + Tensorflow and Multiprocessing in Python. However, this still leaves me with the dilemma of not knowing how to actually "choose" a gpu to operate the process on.

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

My question boils down to: how does one parallelize prediction for one model in Keras across multiple gpus when using Tensorflow as Keras' backend?

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

Additionally I am curious if similar parallelization for prediction is possible with only one gpu.

高度描述或代码示例将不胜感激！

A high level description or code example would be greatly appreciated!

谢谢！

推荐答案

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

I created one simple example to show how to run keras model across multiple gpus. Basically, multiple processes are created and each of process owns a gpu. To specify the gpu id in process, setting env variable CUDA_VISIBLE_DEVICES is a very straightforward way (os.environ["CUDA_VISIBLE_DEVICES"]). Hope this git repo can help you.

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

这篇关于Keras + Tensorflow:对多 GPU 的预测的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！