PyTorch源码解读之torchvision.models

PyTorch框架中有一个非常重要且好用的包:torchvision,该包主要由3个子包组成,分别是:torchvision.datasets、torchvision.models、torchvision.transforms。这3个子包的具体介绍可以参考官网:http://pytorch.org/docs/master/torchvision/index.html。具体代码可以参考github:https://github.com/pytorch/vision/tree/master/torchvision

这篇博客介绍torchvision.models。torchvision.models这个包中包含alexnet、densenet、inception、resnet、squeezenet、vgg等常用的网络结构,并且提供了预训练模型,可以通过简单调用来读取网络结构和预训练模型。

使用例子:

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

这样就导入了resnet50的预训练模型了。如果只需要网络结构,不需要用预训练模型的参数来初始化,那么就是:

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

如果要导入densenet模型也是同样的道理,比如导入densenet169,且不需要是预训练的模型:

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

由于pretrained参数默认是False,所以等价于:

model = torchvision.models.densenet169()

不过为了代码清晰,最好还是加上参数赋值。

接下来以导入resnet50为例介绍具体导入模型时候的源码。运行model = torchvision.models.resnet50(pretrained=True)的时候,是通过models包下的resnet.py脚本进行的,源码如下:

首先是导入必要的库,其中model_zoo是和导入预训练模型相关的包,另外all变量定义了可以从外部import的函数名或类名。这也是前面为什么可以用torchvision.models.resnet50()来调用的原因。model_urls这个字典是预训练模型的下载地址。

import torch.nn as nn
import math
import torch.utils.model_zoo as model_zoo

__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
           'resnet152']

model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}

接下来就是resnet50这个函数了,参数pretrained默认是False。首先model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)是构建网络结构,Bottleneck是另外一个构建bottleneck的类,在ResNet网络结构的构建中有很多重复的子结构,这些子结构就是通过Bottleneck类来构建的,后面会介绍。然后如果参数pretrained是True,那么就会通过model_zoo.py中的load_url函数根据model_urls字典下载或导入相应的预训练模型。最后通过调用model的load_state_dict方法用预训练的模型参数来初始化你构建的网络结构,这个方法就是PyTorch中通用的用一个模型的参数初始化另一个模型的层的操作。load_state_dict方法还有一个重要的参数是strict,该参数默认是True,表示预训练模型的层和你的网络结构层严格对应相等(比如层名和维度)。

def resnet50(pretrained=False, **kwargs):
    """Constructs a ResNet-50 model.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls[
评论 33
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值