pytorch模型保存与加载以及常见问题

一. 模型保存与加载

#多gpu
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2,3,4' #choose
model = TheModelClass(*args, **kwargs)
model = torch.nn.DataParallel(model).cuda()

#加载与训练模型 file.pth.tar
checkpoint = torch.load('file.pth.tar')
model.load_state_dict[checkpoint]

#保存训练模型，保存路径为model_file
torch.save(model.state_dict(), model_file)

二.常见问题

1.Missing key(s) in state_dict: Unexpected key(s) in state_dict:
如果加载的预训练模型之前使用了torch.nn.DataParallel()，而此时的训练并没有使用，则会出现这样的错误。【该问题常出现在"训练模型与当前模型参数不完全一致时，需要update，而update又必须在DataParallel之前" 的这种情况下】
我们需要去掉参数中的前缀 module.

#方法1:
model.load_state_dict({
   k.replace('module.',''):v for k,v in torch.load('myfile.pth').items()})</