【Keras】多GPU训练和模型保存

本文主要是介绍【Keras】多GPU训练和模型保存，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

先看完整示例

import osimport tensorflow as tf
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.applications.inception_v3 import InceptionV3
from keras.applications.mobilenet import MobileNet
from keras import  layers
from keras.models import Model
from keras.optimizers import RMSprop
from keras import backend as K
from keras.applications.vgg16 import VGG16
from keras.models import Sequential
from keras.layers import Conv2D, MaxPool2D, Activation, Dropout, Flatten, Dense
from keras.optimizers import SGD
from keras.preprocessing.image import ImageDataGenerator, img_to_array, load_img
import numpy as np
import keras
from keras.preprocessing.image import ImageDataGenerator
from keras.layers import Conv2D, MaxPool2D, Activation, Dropout, Flatten, Dense, GlobalMaxPooling2D
from keras.preprocessing.image import ImageDataGenerator
from keras import optimizers
from keras.models import Sequential
from keras.layers import Dropout, Flatten, Dense
from keras import initializers
from keras.callbacks import ModelCheckpoint, EarlyStopping
from keras.applications.vgg19 import VGG19
from keras.applications.resnet50 import ResNet50
from keras.applications.mobilenet_v2 import MobileNetV2
from keras.preprocessing.image import ImageDataGeneratorfrom keras.utils import multi_gpu_modeldef training_vis(hist):loss = hist.history['loss']val_loss = hist.history['val_loss']acc = hist.history['acc']val_acc = hist.history['val_acc']# make a figurefig = plt.figure(figsize=(8,4))# subplot lossax1 = fig.add_subplot(121)ax1.plot(loss,label='train_loss')ax1.plot(val_loss,label='val_loss')ax1.set_xlabel('Epochs')ax1.set_ylabel('Loss')ax1.set_title('Loss on Training and Validation Data')ax1.legend()# subplot accax2 = fig.add_subplot(122)ax2.plot(acc,label='train_acc')ax2.plot(val_acc,label='val_acc')ax2.set_xlabel('Epochs')ax2.set_ylabel('Accuracy')ax2.set_title('Accuracy  on Training and Validation Data')ax2.legend()plt.tight_layout()GPU_IDS = '0,1,2,3'
os.environ["CUDA_VISIBLE_DEVICES"] = GPU_IDS
GPU_NUM = len(GPU_IDS.split(','))
train_data_path = r"/root/SMJ/train"
val_data_path = r"/root/SMJ/val"
model_save_path = r"/root/SMJ/checkpoints"
model_name = "resnet50"
clc_num = 2
batch_size = 32
nb_train_samples = 0
nb_validation_samples = 0
for cls_dir in os.listdir(train_data_path):nb_train_samples += len(os.listdir(os.path.join(train_data_path, cls_dir)))
for cls_dir in os.listdir(val_data_path):nb_validation_samples += len(os.listdir(os.path.join(val_data_path, cls_dir)))if GPU_NUM <= 1:print("Training with 1 GPU")pre_trained_model = ResNet50(include_top=False, weights='imagenet', input_shape=(224, 224, 3))x = GlobalMaxPooling2D()(pre_trained_model.output)x = Dense(256, activation='relu')(x)x = Dropout(0.5)(x)x = Dense(clc_num, activation='softmax')(x)model = Model(inputs=pre_trained_model.input, outputs=x)
else:print("Training with {} GPU".format(GPU_NUM))with tf.device("/cpu:0"):pre_trained_model = ResNet50(include_top=False, weights='imagenet', input_shape=(224, 224, 3))x = GlobalMaxPooling2D()(pre_trained_model.output)x = Dense(256, activation='relu')(x)x = Dropout(0.5)(x)x = Dense(clc_num, activation='softmax')(x)model = Model(inputs=pre_trained_model.input, outputs=x)model = multi_gpu_model(model, gpus=4)train_datagen = ImageDataGenerator(rescale = 1/255,         # horizontal_flip = True,# vertical_flip = True,# rotation_range = 359,
)
test_datagen = ImageDataGenerator(rescale = 1/255,       
)train_generator = train_datagen.flow_from_directory(train_data_path,target_size=(224,224),batch_size=batch_size,
)test_generator = test_datagen.flow_from_directory(val_data_path,target_size=(224,224),batch_size=batch_size,
)
sgd = optimizers.SGD(lr=0.0001, decay=1e-6, momentum=0.9, nesterov=True)
#for layer in model.layers[:30]:
#    layer.trainable=False
model.compile(optimizer=sgd, loss='categorical_crossentropy', metrics=['accuracy'])# 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)# model_checkpoint = ParallelModelCheckpoint(model, filepath=os.path.join(model_save_path, model_name + "-{epoch:02d}_loss-{loss:.4f}_val_loss-{val_loss:.4f}.h5"), monitor='val_loss', verbose=1, save_best_only=True, save_weights_only=True) # 解决多GPU运行下保存模型报错的问题model_checkpoint = ModelCheckpoint(filepath=os.path.join(model_save_path, model_name + "-{epoch:02d}_loss-{loss:.4f}_val_loss-{val_loss:.4f}.h5"),monitor='val_loss',verbose=1,save_best_only=True,save_weights_only=True,mode='auto',period=1)hist=model.fit_generator(train_generator,steps_per_epoch=nb_train_samples//(batch_size*GPU_NUM),epochs=500,validation_data=test_generator,validation_steps=nb_validation_samples//(batch_size*GPU_NUM),callbacks=[model_checkpoint])

multi_gpu_model

keras.utils.multi_gpu_model(model, gpus=None, cpu_merge=True, cpu_relocation=False)

将模型复制到不同的 GPU 上。

具体来说，该功能实现了单机多 GPU 数据并行性。它的工作原理如下：

将模型的输入分成多个子批次。
在每个子批次上应用模型副本。每个模型副本都在专用 GPU 上执行。
将结果（在 CPU 上）连接成一个大批量。

例如，如果你的 batch_size 是 64，且你使用 gpus=2，那么我们将把输入分为两个 32 个样本的子批次，在 1 个 GPU 上处理 1 个子批次，然后返回完整批次的 64 个处理过的样本。

这实现了多达 8 个 GPU 的准线性加速。

此功能目前仅适用于 TensorFlow 后端。

参数

model: 一个 Keras 模型实例。为了避免OOM错误，该模型可以建立在 CPU 上，详见下面的使用样例。
gpus: 整数 >= 2 或整数列表，创建模型副本的 GPU 数量，或 GPU ID 的列表。
cpu_merge: 一个布尔值，用于标识是否强制合并 CPU 范围内的模型权重。
cpu_relocation: 一个布尔值，用来确定是否在 CPU 的范围内创建模型的权重。如果模型没有在任何一个设备范围内定义，您仍然可以通过激活这个选项来拯救它。

一个 Keras Model 实例，它可以像初始 model 参数一样使用，但它将工作负载分布在多个 GPU 上。

例子

例 1 - 训练在 CPU 上合并权重的模型

import tensorflow as tf
from keras.applications import Xception
from keras.utils import multi_gpu_model
import numpy as npnum_samples = 1000
height = 224
width = 224
num_classes = 1000# 实例化基础模型（或者「模版」模型）。
# 我们推荐在 CPU 设备范围内做此操作，
# 这样模型的权重就会存储在 CPU 内存中。
# 否则它们会存储在 GPU 上，而完全被共享。
with tf.device('/cpu:0'):model = Xception(weights=None,input_shape=(height, width, 3),classes=num_classes)# 复制模型到 8 个 GPU 上。
# 这假设你的机器有 8 个可用 GPU。
parallel_model = multi_gpu_model(model, gpus=8)
parallel_model.compile(loss='categorical_crossentropy',optimizer='rmsprop')# 生成虚拟数据
x = np.random.random((num_samples, height, width, 3))
y = np.random.random((num_samples, num_classes))# 这个 `fit` 调用将分布在 8 个 GPU 上。
# 由于 batch size 是 256, 每个 GPU 将处理 32 个样本。
parallel_model.fit(x, y, epochs=20, batch_size=256)# 通过模版模型存储模型（共享相同权重）：
model.save('my_model.h5')

例 2 - 训练在 CPU 上利用 cpu_relocation 合并权重的模型

..
# 不需要更改模型定义的设备范围：
model = Xception(weights=None, ..)try:parallel_model = multi_gpu_model(model, cpu_relocation=True)print("Training using multiple GPUs..")
except ValueError:parallel_model = modelprint("Training using single GPU or CPU..")
parallel_model.compile(..)
..

例 3 - 训练在 GPU 上合并权重的模型（建议用于 NV-link）

..
# 不需要更改模型定义的设备范围：
model = Xception(weights=None, ..)try:parallel_model = multi_gpu_model(model, cpu_merge=False)print("Training using multiple GPUs..")
except:parallel_model = modelprint("Training using single GPU or CPU..")parallel_model.compile(..)
..

关于模型保存

要保存多 GPU 模型，请通过模板模型（传递给 multi_gpu_model 的参数）调用 .save(fname) 或 .save_weights(fname) 以进行存储，而不是通过 multi_gpu_model 返回的模型。

意思就是直接使用传入方法keras.utils.multi_gpu_model(model, gpus)中的model即可，而不要使用返回的parallel_model，即：

model.save('xxx.h5')

Modelcheckpoint callback 报错

Modelchecpoint callback函数无法调用也是因为保存时调用的是paralleled_model.save()导致的。
解决方法一：在modelcheckpoint里参数加上save_weights_only=True后，只会保存模型权重，但是保存的模型只能在同样数量的GPU上载入，而没办法再单GPU下载入。

解决方法二：自定义Callback函数，将callback函数中的paralle model 置为 template 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 = modelsuper(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)check_point = ParallelModelCheckpoint(single_model ,'best.h5')

载入

（这里是使用了上述解决方法一所保存的模型）多GPU下保存的权重无法在单GPU/CPU下运行，预测时可以创建多GPU模型后先导入多GPU的权重而用template model 进行predict：

model = get_model()
paralleled_model=multi_gpu_model(model,gpus=num_gpu)
paralleled_model.load_weights("weights_multi_gpu.h5") # 此时model也自动载入了权重，可用model进行预测
model.predict()