keras的单机多卡并行计算简单粗暴,可是。。。。。。。
训练一时爽,重载火葬场,哭哭。下面上干货
首先单CPU/GPU模型转多GPU模型:
from keras.utils import multi_gpu_model
model = getmodel()#单卡模型
n_GPUs = 3# >=2
mode_M = multi_gpu_model(model,n_GPUs)# ok,转好了
是不是简单粗暴啊,然而问题来了,当训练完成之后该如何保存呢?注意mode_M在训练后还是并行模型,直接保存将不再适用单卡的载入和使用,必需要在同样的环境下才能读入,因此我们应该:
model.save("model.h5")
这样我们在载入模型时就不会出现某些神奇的问题比如只剩一层模型这样的骚操作。(哭哭)
可是当我们想保存检查点的时候,问题就来了,保存的全是并行模型(再哭哭)。为了解决这个问题可在多卡端将该模型转成单卡模型即可解决。
from keras.utils import multi_gpu_model
model_dir = ""
model = get_model()
n_GPUs = 2
model_M = multi_gpu_model(model,n_GPUs)
model_M.load_weights(model_dir)
model.summary()#验证一下
model.save(model_dir)