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