Keras多GPU训练

2021-11-23 13:51:46 浏览数 (1)

Keras 2.X版本后可以很方便的支持使用多GPU进行训练了，使用多GPU可以提高我们的训练过程，比如加速和解决内存不足问题。

多GPU其实分为两种使用情况：数据并行和设备并行。

数据并行将目标模型在多个设备上各复制一份，并使用每个设备上的复制品处理整个数据集的不同部分数据。Keras在 keras.utils.multi_gpu_model 中提供有内置函数，该函数可以产生任意模型的数据并行版本，最高支持在8片GPU上并行。

数据并行是指将我们的模型放到多个GPU上去跑，来处理数据集的不同部分，Keras的keras.utils.multi_gpu_model支持任意模型的数据并行，最多支持8个GPU。我们大多数时候要用到的都是数据并行，其他需求可以参考这篇博客：Keras多GPU及分布式。

这里就给出数据并行的多GPU训练示例：

代码语言：javascript复制
from keras.utils.training_utils import multi_gpu_model   #导入keras多GPU函数

model = get_model()
parallel_model = multi_gpu_model(model, gpus=2) # 设置使用2个gpu，该句放在模型compile之前
parallel_model.compile(loss='categorical_crossentropy', optimizer=Adam(), metrics=['accuracy'])
hist = parallel_model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs_num, validation_data=(x_test, y_test), verbose=1, callbacks=callbacks)

你还可以指定要哪几个GPU来跑：

代码语言：javascript复制
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "3,5"

使用命令“nvidia-smi”可以查看各GPU的使用情况和序号，上面代码就是指定用序号为3和5的两个GPU来跑训练。

其实这样就可以了，就是这么简单。

但是实际运行中，我还是遇到了一些报错。

报错1：

ValueError: Variable batch_normalization_1/moving_mean/biased already exists, disallowed. Did you mean to set reuse=True in VarScope? Originally defined at:

我使用单GPU训练的时候没有问题，改成多GPU后出现这个问题。这个问题好解决，将Tensorflow升级到1.4即可。

报错2：

TypeError: can't pickle ...(different text at different situation) objects

查找资料后，发现可能源于callbacks.ModelCheckpoint() 并进行多 gpu 并行计算时，使用姿势不对导致callbacks 函数报错。

我在代码中为了保存最优的训练模型，加了这个callback：

代码语言：javascript复制
checkpoint = ModelCheckpoint(filepath='./cifar10_resnet_ckpt.h5', monitor='val_acc', verbose=1,save_best_only=True)

而在改为多GPU训练后，每次回调存储的模型变成了parallel_model，这会导致报错，只需要改成依然保存原本的model即可，所以我们需要改一下：

代码语言：javascript复制
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, filepath='./cifar10_resnet_ckpt.h5', monitor='val_acc', verbose=1, save_best_only=True) # 解决多GPU运行下保存模型报错的问题

其余的不变，也就是改为依然存储原本的model即可。还有其他的改法可以参考这篇博客：[Keras] 使用多 gpu 并行训练并使用 ModelCheckpoint() 可能遇到的问题，思路都是一样的，只是改法不同。

这样就能够成功使用多GPU训练啦。

keras

0 人点赞