keras多显卡训练

深夜虫鸣

于 2018-12-04 10:25:34 发布

阅读量4k

点赞数 2

分类专栏：深度学习实操

本文链接：https://blog.youkuaiyun.com/u010122972/article/details/84784245

版权

在Keras中，默认使用单 GPU 进行训练。若要利用多 GPU，需借助 `multi_gpu_model` 函数。首先导入该函数，然后在模型定义后，使用它指定GPU数量。每个GPU会分配相应的批次数据进行计算，CPU负责合并梯度更新权重。在训练和保存模型时，确保保存原始模型而非多 GPU 模型，以适应不同 GPU 数量的场景。如果使用回调函数，需自定义保存原始模型的逻辑。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

使用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上进行梯度计算。如果batchsize为128，显卡n=2，则每张显卡单独计算128/2=64张图像，然后在CPU上将两张显卡计算得到的梯度进行融合更新，并对模型权重进行更新后再将新模型拷贝到GPU再次训练。
（3）从上面可以看出，进行训练时，仍然在model_parallel上进行：

model_parallel.fit(...) #注意是model_parallel

最低0.47元/天解锁文章