pytorch 模型保存与加载

  一、模型保存有两种形式：保存整体模型（包括模型结构和参数）、只保存模型参数

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()”了