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']
之后,模型即可用于测试与加载训练了