要先利用GPU训练,CPU测试,那么在模型训练时候,是能保存模型的参数而不能保存整个模型,可见保存、加载、测试
💥这里主要讲一点重要的,即在pytorch 1.6的版本中训练模型保存时,不能直接使用
torch.save(state_r, model_out_r_path)
否则,在CPU测试时,由于版本的不兼容会导致出错,正确使用如下:
torch.save(state_g, model_out_g_path, _use_new_zipfile_serialization=False)
而后,在模型加载时候,要限定map_location = cpu的加载:
checkpoint_r = torch.load(opt.model_r,map_location='cpu')
model_rt.load_state_dict(checkpoint_r['model'])
model_rt.eval()
optimizer_rt.load_state_dict(checkpoint_r['optimizer'])
epochs_rt = checkpoint_r['epoch']
GPU下训练的模型即可方便的在CPU环境中测试了
若模型已经训练保存,但是有没有使用_use_new_zipfile_serialization=False来进行约束,那么,可以在pytorch 1.6中直接加载模型,然后再次使用torch.save进行保存为非z