接着上一篇 Google 阵营Tensorflow /Keras为我们提供了大量pre-trained经典卷积网络模型, 这一篇继续介绍Pytorch如何使用Pytorch torchvision.models
中的经典模型进行快速迁移学习。
相比Tensorflow主要多了Facebook团队发布的 SqueezeNet、ShuffleNet、Wide ResNet、MNASNet模型。
torchvision.models
Pytorch一共提供29个经典卷积神经模型
Pytorch官网TORCHVISION.MODELS
迁移学习代码实现
迁移学习十分简单,直接使用torchvision.models.Xnet去掉最后一个modul便可以实现。
import torchvision
base_model=torchvision.models.alexnet(pretrained=False)
model = nn.Sequential(*list(trained_model.children())[:-1],
Flatten(),
nn.Linear(512, 5)
).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criteon = nn.CrossEntropyLoss()
另外,Pytorch十分方便的地方是所有模型的源代码可以直接从官网推荐中进行拷贝
,比如,点击SOURCE,便可以得到Alexnet的源代码和预训练权重文件地址
torchvision.models网址
import torch
import torch.nn as nn
# load_state_dict_from_url加载地址替换为torch.hub
#from .utils import load_state_dict_from_url
try:
from torch.hub import load_state_dict_from_url
except ImportError:
from torch.utils.model_zoo import load_url as load_state_dict_from_url
__all__ = ['AlexNet', 'alexnet']
model_urls = {
'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth',