keras.fit_generator及多卡训练记录

最新推荐文章于 2025-05-14 22:27:35 发布

qq_33116953

最新推荐文章于 2025-05-14 22:27:35 发布

阅读量603

点赞数

文章标签：深度学习 generator

本文链接：https://blog.csdn.net/qq_33116953/article/details/105256924

版权

1.环境问题

使用keras，以tensorflow为背景，tensorflow1.14多卡训练会出错 python3.6

2.代码

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保存