keras使用多GPU训练

曾经天真的我以为加了下面这个就已经使用了多个GPU训练，事实上，它只用了其他卡的显存。

import os
os.environ['CUDA_VISIBLE_DEVICES']='0,1' #指定哪几块GPU

后来经过查找了一波资料后，终于找到了真正用多GPU训练的方法，这个方法也很简单，从上面的基础上再插入一个函数就可以了。

实验条件：

tensorflow 1.13.1

keras 2.2.4

先导入所需要的函数

from keras.utils import multi_gpu_model

就是这个函数，注意：tensorflow 1.14.0里面的这个包有bug

传送门：https://www.tensorflow.org/api_docs/python/tf/keras/utils/multi_gpu_model

然后在model.compile之前用这个函数，比如下面这个例子我写了给resnet 50的model

pre_model = ResNet50((512,512,3))
model=multi_gpu_model(pre_model, gpus=2) #就插入到这里
opt = optimizers.Adam(lr=0.001)
model.compile(optimizer = opt, loss = "binary_crossentropy", metrics = ["accuracy"])

然后就可以使用了

 

 如果你要用callback，要稍微改一下Checkpoint()的使用方法

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='./trained_model/previous.h5', monitor='loss',  save_weights_only=True, verbose=1, save_best_only=True)

这样就可以愉快地使用多GPU了。

转载于:https://www.cnblogs.com/yexionglin/p/11459696.html