一、模型保存有两种形式:保存整体模型(包括模型结构和参数)、只保存模型参数
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()”了