使用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,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,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儲存

以上這篇keras多顯示卡訓練方式就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。