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)