Pytorch模型保存机制

在训练中，模型保存至关重要，pytorch中主要有以下两种保存机制：

🅰保存整个模型，即模型+参数

🅱只保存训练参数

保存整个模型

# 保存
torch.save(model, './model.pth')
# 加载
model = torch.load('./model.pth')

只保存参数

该方法的好处是可以用于后续模型的CPU测试以及加载继续训练

1️⃣只保留模型参数

# 保存
torch.save(model.state_dict(), '\model.pth')
# 加载
model = MyModelClass()  
model.load_state_dict(torch.load('\model.pth'))

因为该方法只保存了模型参数，所以，在加载参数之前，要现有一个现有的模型，相当于：得现有瓶子(模型），才能往里面灌水（参数）

2️⃣保存模型参数及其他信息

def checkpoint(epoch):
    model_out_r_path = join(opt.save_models, 'LapSRN_r_epoch_{}.pt'.format(epoch))
    model_out_g_path = join(opt.save_models, 'LapSRN_g_epoch_{}.pt'.format(epoch))
    state_r = {'model': model_r.state_dict(), 'optimizer': optimizer_r.state_dict(), 'epoch':epoch, 'lr':lr_r}
    state_g = {'model': model_g.state_dict(), 'optimizer': optimizer_g.state_dict(), 'epoch':epoch, 'lr':lr_g}
    torch.save(state_r, model_out_r_path)
    torch.save(state_g, model_out_g_path)
    print("Checkpoint saved to {} and {}".format(model_out_r_path, model_out_g_path))

若要保存模型其他信息，则要利用一个dict来存储而后在进行保存

在模型及相关参数加载时：

model_rt = LapSRN_r().cuda()		# 先加载模型
model_gt = LapSRN_g().cuda()
optimizer_rt = optim.Adagrad(model_rt.parameters(), lr=1e-3, weight_decay=1e-5)# 加载优化器
optimizer_gt = optim.Adagrad(model_gt.parameters(), lr=1e-3, weight_decay=1e-5)

checkpoint_r = torch.load(model_out_r_path)
model_rt.load_state_dict(checkpoint_r['model'])
model_rt.eval()
optimizer_rt.load_state_dict(checkpoint_r['optimizer'])
epochs_rt = checkpoint_r['epoch']

checkpoint_g = torch.load(model_out_g_path)
model_gt.load_state_dict(checkpoint_g['model'])
model_gt.eval()
optimizer_gt.load_state_dict(checkpoint_g['optimizer'])
epochs_gt = checkpoint_g['epoch']

之后，模型即可用于测试与加载训练了