PyTorch模型的保存與加載方法實例

更新時間：2022年09月23日 09:53:59 作者：長命百歲?

Pytorch保存模型其實非常簡單,下面這篇文章主要給大家介紹了關于PyTorch模型的保存與加載的相關資料,文中通過實例代碼介紹的非常詳細,需要的朋友可以參考下

模型的保存與加載

首先，需要導入兩個包

import torch
import torchvision.models as models

保存和加載模型參數

PyTorch模型將學習到的參數存儲在一個內部狀態(tài)字典中，叫做state_dict。這可以通過torch.save方法來實現(xiàn)。
我們導入預訓練好的VGG16模型，并將其保存。我們將state_dict字典保存在model_weights.pth文件中。

model = models.vgg16(pretrained=True)
torch.save(model.state_dict(), 'model_weights.pth')

想要加載模型參數，我們需要創(chuàng)建一個和原模型一樣的實例，然后通過load_state_dict()方法來加載模型參數

創(chuàng)建一個VGG16模型實例（未經過預訓練的）
加載本地參數

model = models.vgg16() # we do not specify pretrained=True, i.e. do not load default weights
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

注意：在進行測試前，如果模型中有dropout層和batch normalization層的話，一定要使用model.eval()將模型轉到測試模式。

在train模式下，dropout網絡層會按照設定的參數p設置保留激活單元的概率（保留概率=p); batchnorm層會繼續(xù)計算數據的mean和var等參數并更新。
在val模式下，dropout層會讓所有的激活單元都通過，而batchnorm層會停止計算和更新mean和var，直接使用在訓練階段已經學出的mean和var值