一、模型保存有两种形式：保存整体模型（包括模型结构和参数）、只保存模型参数
import torch device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 保存整体模型 output_dir = 'checkpoint.ckp' model_to_save = model.module if hasattr(model, "module") else model torch.save(model_to_save, output_dir) # 加载整体模型 model = torch.load(output_dir, map_location=device) # ========================================================================= # 只保存模型参数 state_dict torch.save(model.state_dict(), output_dir) # 只加载模型参数 state_dict model.load_state_dict(torch.load(output_dir, map_location=device))
二、当训练的代码中使用了“torch.nn.DataParallel()”，这个命令是将网络在多块gpu中进行训练然后合并，这时采用上面“只保存模型参数”的方式时，保存的参数key中会在最前面多一个module.
解决方式有三个：
1.加载模型时去掉key中的module.
model.load_state_dict({k.replace('module.',''):v for k,v in torch.load('checkpoint.pth').items()})
2.加载时也用“torch.nn.DataParallel()”
if cuda: g_model = torch.nn.DataParallel(g_model) cudnn.benchmark = True g_model = g_model.cuda() if os.path.exists(model_path): print('Loading weights into state dict...') model_dict = g_model.state_dict() pretrained_dict = torch.load(model_path, map_location=device) g_model.load_state_dict(pretrained_dict) print('Finished!')
3.只有一个GPU就没有必要使用“torch.nn.DataParallel()”了
