如何使用Keras实现多显卡并行训练？

本文共计880个文字，预计阅读时间需要4分钟。

使用Keras进行训练时，默认使用单显卡。已设置环境变量`CUDA_VISIBLE_DEVICES`为两块显卡，确保只占用了显存。再设置`tf.GPUOptions(allow_growth=True)`后，可以看到仅使用了第一块显卡。

使用keras进行训练，默认使用单显卡，即使设置了os.environ['CUDA_VISIBLE_DEVICES']为两张显卡，也只是占满了显存，再设置tf.GPUOptions(allow_growth=True)之后可以清楚看到，只占用了第一张显卡，第二张显卡完全没用。

要使用多张显卡，需要按如下步骤：

（1）import multi_gpu_model函数：from keras.utils import multi_gpu_model

（2）在定义好model之后，使用multi_gpu_model设置模型由几张显卡训练，如下：

model=Model(...) #定义模型结构 model_parallel=multi_gpu_model(model,gpu=n) #使用几张显卡n等于几 model_parallel.compile(...) #注意是model_parallel，不是model

通过以上代码，model将作为CPU上的原始模型，而model_parallel将作为拷贝模型被复制到各个GPU上进行梯度计算。