1. 保存整个模型

torch.save(model, 'model.pkl')
model = torch.load('model.pkl')

2.仅保存和加载模型参数(推荐使用)

torch.save(model_object.state_dict(), 'params.pkl')
model.load_state_dict(torch.load('params.pkl'))

注意在多卡训练的时候，会多保留一个module
因此我们可以设置

if torch.cuda.device_count() > 1:  #多卡训练torch.save(model.module.state_dict(),'params.pkl')
else:torch.save(model.state_dict(),'params.pkl')

3. 保存其他参数到模型中，比如optimizer,epoch 等

torch.save({'epoch':edix, 'state_dict':model.state_dict(),  \
'optimizer':optimizer.state_dict()}, \
f'./saved_models/{opt.exp_name}/epoch_{edix+1}.pth')

model.load_state_dict(torch.load(opt.saved_model)['state_dict'], strict=False)model.load_state_dict(torch.load(opt.saved_model)['epoch'], strict=False)
model.load_state_dict(torch.load(opt.saved_model)['optimizer'], strict=False)