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)