多GPU训练可以在一定程度上提高训练的效率，keras在以tensorflow为底层的时候，可以方便的调用多GPU进行训练。其训练的原理是将一个神经网络在多块GPU中进行初始化，并且将一个batch的数据分配到不同的GPU中，反向传播结束后，再使用CPU进行合并更新参数。举个例子，同时使用两块GPU，batchsize是16的话，keras会将8个给1号GPU，另外8个给2号GPU，训练结束后将二者拼接为一个16的batch并进行反向传播更新参数。
通过以上的原理我们知道，当batchsize是1的时候，比如在三维数据上，使用多个GPU训练实际上是没有提速效果的。
第一步，就是多GPU的选取，比较直接的方式就是通过更改环境变量来选择使用哪个GPU。
import os os.environ["CUDA_VISIBLE_DEVICES"] = "2,3" 1 2 没有指定keras默认会选择所有的GPU，如果你并没有配置多GPU的训练，那么实际上keras占了所有的GPU但是实际上只在一个GPU上运行。
第二步，就是调用多GPU函数，multi_gpu_model，第一个参数是你的模型，第二个是使用GPU的数量。model泛指一个keras的模型。
from keras.callbacks import TensorBoard, ModelCheckpoint
from keras.utils import multi_gpu_model # 导入keras多卡函数
class ParallelModelCheckpoints(ModelCheckpoint): # 在保存模型时,由于存在两个模型,所以需要指定model,\
# 继承ModelCheckpoint,重写init()
def __init__(self,
model, # 需要保存的模型
filepath='./log/epoch-{epoch:02d}_loss-{loss:.4f}_acc-{val_acc:.4f}_lr-{lr:.5f}.h5',
monitor='val_acc',
verbose=1,
save_best_only=True,
save_weights_only=False,
mode='auto',
period=1):
self.single_model = model
super(ParallelModelCheckpoints, self).__init__(filepath, monitor, verbose,save_best_only, save_weights_only, mode, period)
def set_model(self, model):
super(ParallelModelCheckpoints, self).set_model(self.single_model)
# 首先在cpu上创建原来的模型
with tf.device('/cpu:0'):
model = MobileNet(...)
# 创建多卡模型
parallel_model = multi_gpu_model(model, gpus=4) # 其中 4 是gpu的数量
parallel_model.load_weights(h5_path, by_name=True) # 继续训练的时候导入参数是用的parallel_model模型,而不是model
parallel_model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
model_checkpoint = ParallelModelCheckpoints(model) # 设置需要保存h5的模型
print("Start training the model")
# 然后就可以训练了
training_history = parallel_model.fit_generator(
train_generator,
steps_per_epoch=step_size_train,
validation_data=validation_generator,
validation_steps=step_size_valid,
epochs=epoch_list[-1],
verbose=1,
callbacks=[TensorBoard(log_dir='./tb'), model_checkpoint, stepDecayLR])
print("Model training finished")
