pytorch自發(fā)布以來，由于其便捷性，贏得了越來越多人的喜愛。

Pytorch有很多方便易用的包，今天要談的是torchvision包，它包括3個(gè)子包，分別是： torchvison.datasets ，torchvision.models ，torchvision.transforms ，分別是預(yù)定義好的數(shù)據(jù)集（比如MNIST、CIFAR10等）、預(yù)定義好的經(jīng)典網(wǎng)絡(luò)結(jié)構(gòu)（比如AlexNet、VGG、ResNet等）和預(yù)定義好的數(shù)據(jù)增強(qiáng)方法（比如Resize、ToTensor等）。這些方法可以直接調(diào)用，簡(jiǎn)化我們建模的過程，也可以作為我們學(xué)習(xí)或構(gòu)建新的模型的參考。

本文，我們講述的是models，且只談模型的加載。models這個(gè)包中包含alexnet、densenet、inception、resnet、squeezenet、vgg等常用的網(wǎng)絡(luò)結(jié)構(gòu)，并且提供了預(yù)訓(xùn)練模型，可以通過簡(jiǎn)單調(diào)用來讀取網(wǎng)絡(luò)結(jié)構(gòu)和預(yù)訓(xùn)練模型。

模型地址：https://github.com/pytorch/vision/tree/master/torchvision/models

官方文檔：https://pytorch.org/docs/master/torchvision/models.html

我將加載的方法簡(jiǎn)單總結(jié)為以下四種：

1.直接加載預(yù)訓(xùn)練模型

import torchvision.models as models resnet50 = models.resnet50(pretrained=True)

?
這樣就導(dǎo)入了resnet50的預(yù)訓(xùn)練模型了。

如果只需要網(wǎng)絡(luò)結(jié)構(gòu)，不需要用預(yù)訓(xùn)練模型的參數(shù)來初始化，那么就是：

model =torchvision.models.resnet50(pretrained=False)

或者把resnet復(fù)制到自己的目錄下，新建個(gè)model文件夾

可以參考下面的貓狗大戰(zhàn)入門算法入門

https://github.com/JackwithWilshere/Kaggle-Dogs_vs_Cats_PyTorch

2.修改某一層
?以resnet為例，默認(rèn)的是ImageNet的1000類，比如我們要做二分類，分類貓和狗

resnet.fc = nn.Linear(2048, 2)?

resnet 第一層卷積的卷積核是7，我們可能想改成5，那么可以通過以下方法修改：

#未經(jīng)試驗(yàn)，修改需要有理論依據(jù)，計(jì)算featuremap維度使之匹配。

resnet.conv1 = nn.Conv2d(3, 64,kernel_size=5, stride=2, padding=3, bias=False)

3.加載部分預(yù)訓(xùn)練模型
對(duì)于具體的任務(wù)，很難保證模型和公開的模型完全一樣，但是預(yù)訓(xùn)練模型的參數(shù)確實(shí)有助于提高訓(xùn)練的準(zhǔn)確率，為了結(jié)合二者的優(yōu)點(diǎn)，就需要我們加載部分預(yù)訓(xùn)練模型。

#加載model，model是自己定義好的模型 resnet50 = models.resnet50(pretrained=True)? model =Net(...)?#讀取參數(shù)? pretrained_dict =resnet50.state_dict()? model_dict = model.state_dict()?#將pretrained_dict里不屬于model_dict的鍵剔除掉? pretrained_dict = ?{k: v for k, v in pretrained_dict.items() if k in model_dict}?# 更新現(xiàn)有的model_dict? model_dict.update(pretrained_dict)?# 加載我們真正需要的state_dict? model.load_state_dict(model_dict) ? 4. 加載自己的模型 其實(shí)這個(gè)是保存和恢復(fù)模型，比如我們訓(xùn)練好的模型保存，然后加載用于測(cè)試。方法一(推薦)：第一種方法也是官方推薦的方法，只保存和恢復(fù)模型中的參數(shù)。使用這種方法，我們需要自己導(dǎo)入模型的結(jié)構(gòu)信息。（1）保存torch.save(model.state_dict(), PATH)#example torch.save(resnet50.state_dict(),'ckp/model.pth') ? ? （2）恢復(fù)model = ModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH))#example resnet=resnet50(pretrained=True) resnet.load_state_dict(torch.load('ckp/model.pth'))

方法二：

使用這種方法，將會(huì)保存模型的參數(shù)和結(jié)構(gòu)信息。

(1)保存

torch.save (model, PATH)

(2)恢復(fù)

model = torch.load(PATH)

參考資料：

1. https://zhuanlan.zhihu.com/p/25980324

2. http://www.pytorchtutorial.com/pytorch-note5-save-and-restore-models/

3.https://blog.csdn.net/weixin_41278720/article/details/80759933?
?

總結(jié)

以上是生活随笔為你收集整理的PyTorch【torchvision】的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。