如何使用keras多显卡进行训练 使用keras多显卡训练方式代码示例

如何使用keras多显卡进行训练？本篇文章小编给大家分享一下使用keras多显卡训练方式代码示例，小编觉得挺不错的，现在分享给大家供大家参考，有需要的小伙伴们可以来看看。

使用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

（4）保存模型时，model_parallel保存了训练时显卡数量的信息，所以如果直接保存model_parallel的话，只能将模型设置为相同数量的显卡调用，否则训练的模型将不能调用。因此，为了之后的调用方便，只保存CPU上的模型，即model:

model.save(...) #注意是model，不是model_parallel

如果用到了callback函数，则默认保存的也是model_parallel（因为训练函数是针对model_parallel的），所以要用回调函数保存model的话需要自己对回调函数进行定义：

class OwnCheckpoint(keras.callbacks.Callback):
 def __init__(self,model):
  self.model_to_save=model
 def on_epoch_end(self,epoch,logs=None): #这里logs必须写
  self.model_to_save.save('model_advanced/model_%d.h5' % epoch)

定以后具体使用如下：

checkpoint=OwnCheckpoint(model)

model_parallel.fit_generator(...,callbacks=[checkpoint])

这样就没问题了！