keras多顯卡訓練方式

使用keras進行訓練，默認使用單顯卡，即使設置了os.environ['CUDA_VISIBLE_DEVICES']為兩張顯卡，也只是佔滿了顯存，再設置tf.GPUOptions(allow_growth=True)之後可以清楚看到，只佔用了第一張顯卡，第二張顯卡完全沒用。

要使用多張顯卡，需要按如下步驟：

（1）import multi_gpu_model函數：from keras.utils import multi_gpu_model

（2）在定義好model之後，使用multi_gpu_model設置模型由幾張顯卡訓練，如下：

  model=Model(...) #定義模型結構  model_parallel=multi_gpu_model(model,gpu=n) #使用幾張顯卡n等於幾  model_parallel.compile(...) #注意是model_parallel，不是model

通過以上代碼，model將作為CPU上的原始模型，而model_parallel將作為拷貝模型被複制到各個GPU上進行梯度計算。如果batchsize為128，顯卡n=2，則每張顯卡單獨計算128/2=64張圖像，然後在CPU上將兩張顯卡計算得到的梯度進行融合更新，並對模型權重進行更新後再將新模型拷貝到GPU再次訓練。

（3）從上面可以看出，進行訓練時，仍然在model_parallel上進行：

model_parallel.fit(...) #注意是model_parallel

（4）保存模型時，model_parallel保存了訓練時顯卡數量的信息，所以如果直接保存model_parallel的話，只能將模型設置為相同數量的顯卡調用，否則訓練的模型將不能調用。因此，為了之後的調用方便，只保存CPU上的模型，即model:

model.save(...) #注意是model，不是model_parallel

如果用到了callback函數，則默認保存的也是model_parallel（因為訓練函數是針對model_parallel的），所以要用回調函數保存model的話需要自己對回調函數進行定義：

  class OwnCheckpoint(keras.callbacks.Callback):   def __init__(self,model):    self.model_to_save=model   def on_epoch_end(self,epoch,logs=None): #這裡logs必須寫    self.model_to_save.save('model_advanced/model_%d.h5' % epoch)

定以後具體使用如下：

checkpoint=OwnCheckpoint(model)
 model_parallel.fit_generator(...,callbacks=[checkpoint])

這樣就沒問題了！

補充知識：keras.fit_generator及多卡訓練記錄

1.環境問題

使用keras，以tensorflow為背景，tensorflow1.14多卡訓練會出錯 python3.6

2.代碼

2.1

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 os.environ['CUDA_VISIBLE_DEVICES'] = '4,5'

2.2 自定義generator函數

  def img_image_generator(path_img, path_lab, batch_size, data_list):   while True:   # 'train_list.csv'   file_list = pd.read_csv(data_list, sep=',',usecols=[1]).values.tolist()   file_list = [i[0] for i in file_list]   cnt = 0   X = []   Y1 = []   for file_i in file_list:   x = cv2.imread(path_img+'/'+file_i, cv2.IMREAD_GRAYSCALE)   x = x.astype('float32')   x /= 255.   y = cv2.imread(path_lab+'/'+file_i, cv2.IMREAD_GRAYSCALE)   y = y.astype('float32')   y /= 255.   X.append(x.reshape(256, 256, 1))   Y1.append(y.reshape(256, 256, 1))   cnt += 1   if cnt == batch_size:   cnt = 0   yield (np.array(X), [np.array(Y1), np.array(Y1)])   X = []   Y1 = []

2.3 函數調用及訓練

   generator_train = img_image_generator(path1, path2, 4, pathcsv_train)   generator_test= img_image_generator(path1, path2, 4, pathcsv_test)   model.fit_generator(generator_train, steps_per_epoch=237*2, epochs=50, callbacks=callbacks_list, validation_data=generator_test, validation_steps=60*2)

3. 多卡訓練

3.1 複製model

model_parallel = multi_gpu_model(model, gpus=2)

3.2 checkpoint 定義

  class ParallelModelCheckpoint(ModelCheckpoint):    def __init__(self, model, filepath, monitor='val_out_final_score', 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=model, filepath=filepath, monitor='val_loss',verbose=1, save_best_only=True, mode='min')

3.3 注意的問題

保存模型是時候需要使用以原來的模型保存，不能使用model_parallel保存