快捷導(dǎo)航

pytorch 加載(.pth)格式的模型實(shí)例

更新時間：2019年08月20日 11:15:19 作者：z智慧

今天小編就為大家分享一篇pytorch 加載(.pth)格式的模型實(shí)例，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧

有一些非常流行的網(wǎng)絡(luò)如 resnet、squeezenet、densenet等在pytorch里面都有，包括網(wǎng)絡(luò)結(jié)構(gòu)和訓(xùn)練好的模型。

pytorch自帶模型網(wǎng)址：https://pytorch-cn.readthedocs.io/zh/latest/torchvision/torchvision-models/

按官網(wǎng)加載預(yù)訓(xùn)練好的模型：

import torchvision.models as models
 
# pretrained=True就可以使用預(yù)訓(xùn)練的模型
resnet18 = models.resnet18(pretrained=True)
print(resnet18)

報(bào)錯如下：

requests.exceptions.ConnectionError: ('Connection aborted.', TimeoutError(10060, '由于連接方在一段時間后沒有正確答復(fù)或連接的主機(jī)沒有反應(yīng)，連接嘗試失敗。', None, 10060, None))

主要是因?yàn)榇a會去遠(yuǎn)端下載模型的參數(shù)，而國內(nèi)的網(wǎng)一般連接不上，這是我們需要手動去下載你要的預(yù)訓(xùn)練網(wǎng)絡(luò)。

通過地址下載，地址有兩種獲取方式：

1.從報(bào)錯里面獲取，上述代碼運(yùn)行時會出現(xiàn)這樣一行信息：

Downloading: "https://download.pytorch.org/models/resnet18-5c106cde.pth" to C:\Users\Luo/.torch\models\resnet18-5c106cde.pth

復(fù)制這個網(wǎng)址到瀏覽器，有可能打不開，去掉https://，直接輸入download.pytorch.org/models/resnet18-5c106cde.pth就可以下載了。

2.從pytorch的github下找模型的地址：https://github.com/pytorch/vision/tree/master/torchvision/models

找到對應(yīng)模型名稱點(diǎn)進(jìn)去找地址

下載好后自行保存，我是直接存在pytorch models里面

接下來就是運(yùn)行這個.pth文件。首先要判斷是保存的整個網(wǎng)絡(luò)結(jié)構(gòu)加參數(shù)呢，還是只保存了參數(shù)，可以測試一下。這是我的模型是squeezenet1_1，你可以測試自己下載的模型

import torch
pthfile = r'E:\anaconda\app\envs\luo\Lib\site-packages\torchvision\models\squeezenet1_1.pth'
net = torch.load(pthfile)
print(net)

結(jié)果為

很明顯就是只保存了參數(shù)，這是我們要換個方法加載模型

import torch
import torchvision.models as models
 
# pretrained=True就可以使用預(yù)訓(xùn)練的模型
net = models.squeezenet1_1(pretrained=False)
pthfile = r'E:\anaconda\app\envs\luo\Lib\site-packages\torchvision\models\squeezenet1_1.pth'
net.load_state_dict(torch.load(pthfile))
print(net)

結(jié)果;