本人用pytorch框架在两块GPU上并行训练了一个神经网络模型,并将训练的不同阶段的结果保存起来,以便用于模型集成。
虽然模型是在GPU上训练的,但是在服务器上部署的时候只需用CPU就可以进行模型推断。但在实际应用中,却出现如下报错信息:
RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cpu
解决的思路是在服务器上先将模型加载进来,然后用一种新的方式重新保存。见下面的代码段:
model_list = ['model_1.tar','model_2.tar','model_3.tar','model_4.tar','model_5.tar']
model_path = './sel_models/'
# set model
device = torch.device('cpu')
for model_name in model_list:
model = Net(num_classes=num_classes, num_channels=num_channels).to(device, dtype=torch.float)
model = nn.DataParallel(model)
# load trained model
checkpoint = torch.load(os.path.join(model_path, model_name), map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
del checkpoint
model = model.to(device, dtype=torch.float)
cpu_model_path = './model_for_cpu/'
if not os.path.exists(cpu_model_path):
os.mkdir(cpu_model_path)
torch.save({'model_state_dict': model.module.state_dict()},os.path.join(cpu_model_path, model_name))
这里需要注意的是 必须要加上 model = nn.DataParallel(model), 因为模型是在双GPU上并行训练的,不加这句话模型加载就会出错。另外重新保存的时候一定要加上 'module', 即 model.module.state_dict(), 而不是model.state_dict(),这也是解决这个错误的关键。
本文参考: