Keras使用多GPU训练模型、保存、加载

多GPU训练可以在一定程度上提高训练的效率，Keras在以TensorFlow为底层的时候，可以方便的调用多GPU进行训练。其训练的原理是将一个神经网络在多块GPU中进行初始化，并且将一个batch的数据分配到不同的GPU中，反向传播结束后，再使用CPU进行合并更新参数。举个例子，同时使用两块GPU，batchsize是16的话，Keras会将8个给1号GPU，另外8个给2号GPU，训练结束后将二者拼接为一个16的batch并进行反向传播更新参数。

1. 可用的GPU资源

通过环境变量来设置可用的GPU资源，有以下两种办法：

1）在代码里面设置

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

2）通过命令行设置

CUDA_VISIBLE_DEVICES=0,1 python -W ignore LIDC-IDRI_train.py

没有指定的话，Keras默认会占用所有的GPU，但如果没有配置多GPU训练，Keras占实际上只在一个GPU上训练。

2. 使用多GPU训练

下面这个例子来自于官方文档，https://keras-cn.readthedocs.io/en/latest/utils/#multi_gpu_model

import tensorflow as tf
from keras.applications import Xception
from keras.utils import multi_gpu_model
import numpy as np

num_samples = 1000
height = 224
width = 224
num_classes = 1000

# Instantiate the base model
# (here, we do it on CPU, which is optional).
with tf.device('/cpu:0'):
    model = Xception(weights=None,
                     input_shape=(height, width, 3),
                     classes=num_classes)

# Replicates the model on 8 GPUs.
# This assumes that your machine has 8 available GPUs.
parallel_model = multi_gpu_model(model, gpus=8)
parallel_model.compile(loss='categorical_crossentropy',
                       optimizer='rmsprop')

# Generate dummy data.
x = np.random.random((num_samples, height, width, 3))
y = np.random.random((num_samples, num_classes))

# This `fit` call will be distributed on 8 GPUs.
# Since the batch size is 256, each GPU will process 32 samples.
parallel_model.fit(x, y, epochs=20, batch_size=256)

3. 模型保存

当需要使用Keras的checkpoint的功能的时候，使用如下函数。注意：初始化的参数model，就是单个的model，而不是用于多GPU训练的parallel_model。

class ParallelModelCheckpoint(ModelCheckpoint):
    def __init__(self,model,filepath, monitor='val_loss', verbose=0,
                 save_best_only=False, save_weights_only=False,
                 mode='auto', period=1):
        self.single_model = model
        super(ParallelModelCheckpoint,self).__init__(filepath, monitor, verbose,save_best_only, save_weights_only,mode, period)

    def set_model(self, model):
        super(ParallelModelCheckpoint,self).set_model(self.single_model)

调用如下：

checkpoint = ParallelModelCheckpoint(model, save5_path, monitor='val_auc', verbose=1, save_best_only=True, save_weights_only=False, mode='max')

callbacks = [check_point, early_stopping, tbCallBack]

model.fit(x_train, y_train,
          batch_size=config.batch_size,
          epochs=config.nb_epoch, 
          verbose=config.verbose, 
          shuffle=True,
          validation_data=(x_valid, y_valid),
          callbacks=callbacks)

4. 模型加载与测试

如果模型保存时，save_weights_only=False，则可以直接使用load_model来加载模型。

from keras.models import load_model

model = load_model(model_path, custom_objects={'bce_dice_loss': bce_dice_loss, 'dice_coef_loss': dice_coef_loss,'dice_coef': dice_coef, 'mean_iou': mean_iou})

p_test = model.predict(x_test, batch_size=config.batch_size, verbose=config.verbose)
eva = model.evaluate(x_test, y_test, batch_size=config.batch_size, verbose=config.verbose)