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了。