文章目錄

• 1. 保存整個模型
• 2.僅保存和加載模型參數(推薦使用)
• 3. 保存其他參數到模型中，比如optimizer,epoch 等

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) 創作挑戰賽新人創作獎勵來咯，堅持創作打卡瓜分現金大獎

總結

以上是生活随笔為你收集整理的pytorch: 在训练中保存模型,加载模型的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。