快捷導(dǎo)航

keras多顯卡訓(xùn)練方式

更新時(shí)間：2020年06月10日 15:32:16 作者：深夜蟲鳴

這篇文章主要介紹了keras多顯卡訓(xùn)練方式，具有很好的參考價(jià)值，希望對大家有所幫助。一起跟隨小編過來看看吧

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

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

（1）import multi_gpu_model函數(shù)：from keras.utils import multi_gpu_model

（2）在定義好model之后，使用multi_gpu_model設(shè)置模型由幾張顯卡訓(xùn)練，如下：

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

通過以上代碼，model將作為CPU上的原始模型，而model_parallel將作為拷貝模型被復(fù)制到各個(gè)GPU上進(jìn)行梯度計(jì)算。如果batchsize為128，顯卡n=2，則每張顯卡單獨(dú)計(jì)算128/2=64張圖像，然后在CPU上將兩張顯卡計(jì)算得到的梯度進(jìn)行融合更新，并對模型權(quán)重進(jìn)行更新后再將新模型拷貝到GPU再次訓(xùn)練。

（3）從上面可以看出，進(jìn)行訓(xùn)練時(shí)，仍然在model_parallel上進(jìn)行：

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

（4）保存模型時(shí)，model_parallel保存了訓(xùn)練時(shí)顯卡數(shù)量的信息，所以如果直接保存model_parallel的話，只能將模型設(shè)置為相同數(shù)量的顯卡調(diào)用，否則訓(xùn)練的模型將不能調(diào)用。因此，為了之后的調(diào)用方便，只保存CPU上的模型，即model:

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

如果用到了callback函數(shù)，則默認(rèn)保存的也是model_parallel（因?yàn)橛?xùn)練函數(shù)是針對model_parallel的），所以要用回調(diào)函數(shù)保存model的話需要自己對回調(diào)函數(shù)進(jìn)行定義：

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])

這樣就沒問題了！

補(bǔ)充知識：keras.fit_generator及多卡訓(xùn)練記錄

1.環(huán)境問題

使用keras，以tensorflow為背景，tensorflow1.14多卡訓(xùn)練會出錯(cuò) python3.6

2.代碼

2.1

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

2.2 自定義generator函數(shù)

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 函數(shù)調(diào)用及訓(xùn)練

 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. 多卡訓(xùn)練

3.1 復(fù)制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)

使用