PyTorch源碼解讀之torchvision.models

https://blog.csdn.net/u014380165/article/details/79119664

1. 構建預訓練模型：pretrained=True

import torchvision.models as models

resnet18 = models.resnet18(pretrained=True)
vgg16 = models.vgg16(pretrained=True)
alexnet = models.alexnet(pretrained=True)
squeezenet = models.squeezenet1_0(pretrained=True)

預訓練模型的期望輸入: RGB圖像的mini-batch: (batch_size, 3, H, W)。

其中H, W>=224。

圖像的像素值必須在范圍[0,1]間，

並且用均值mean=[0.485, 0.456, 0.406]和方差std=[0.229, 0.224, 0.225]進行歸一化。

如果只需要網絡結構，不需要用預訓練模型的參數來初始化，用pretrained=False

model = torchvision.models.densenet169(pretrained=False)
# 等價於：
model = torchvision.models.densenet169()

例子：

import torchvision.models as models

vgg16 = models.vgg16(pretrained = True) # 獲取訓練好的VGG16模型
pretrained_dict = vgg16.state_dict() # 返回包含模塊所有狀態的字典，包括參數和緩存

 

2.  源碼分析：以resnet50為例。

運行model = resnet50(pretrained=True)時，通過models包下的resnet.py腳本進行。

resnet.py腳本源碼：

1) 首先導入必要的庫；

model_zoo是和導入預訓練模型相關的包；

all變量定義了可以從外部導入的函數名或類名；

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',
}

2) resnet50函數：

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['resnet50']))
    return model

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

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
    return model

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

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
    return model

 

- 構建網絡結構：model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)；

　　- Bottleneck是另外一個構建bottleneck的類，

　　- 在ResNet網絡結構的構件中有很多重復的子結構，

　　- 這些子結構就是通過Bottleneck類來構建的。 

-  model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))

       - pretrained=True，會通過model_zoo.py中的load_url函數根據model_urls字典下載或導入相應的預訓練模型。

       - 通過調用model的load_state_dict()用預訓練的模型參數來初始化你構建的網絡結構，

       　　- 這個方法就是PyTorch中通用的用一個模型的參數初始化另一個模型的層的操作。

　　　　- load_state_dict()還有一個重要的參數是strict，默認是True

　　　　- 表示預訓練模型的層和你的網絡結構層嚴格對應相等（比如層名和維度）。

- 其他resnet18、resnet101、resnet152的區別：

　　a. 構建網絡結構的時候block參數不同：

　　　　a1. resnet18:[2,2,2,2]; resnet101:[3,4,23,3];

　　b. 調用的block類不一樣:

　　　　b1. resnet50、resnet101、resnet152調用Bottleneck類；

　　　　b2. resnet18、resnet34調用BasicBlock

　　　　b3. 兩者區別：在residual結果中卷積層的數量不同，與網絡結構相關。

　　c. 如果下載預訓練模型，model_urls字典的鍵不一樣，對應不同的預訓練模型。

3. 如何構建網絡、如何導入預訓練模型。

class ResNet(nn.Module):

    def __init__(self, block, layers, num_classes=1000):
        self.inplanes = 64
        super(ResNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x

1）構建ResNet網絡是通過ResNet類進行的

- 繼承Pytorch中網絡的基類：torch.nn.Module；

- 重寫初始化__init__()和forward()方法【層的連接順序】；

- 在類中定義其他私有方法用來模塊化一些操作；

　　- 例如，_make_layer()：構建ResNet網絡中的4個blocks。

　　　　- 第一個輸入block是Bottleneck或BasicBlock類；

　　　　- 第二個輸入是該blocks的輸出channel；

　　　　- 第三個輸入是每個blocks中包含多少個residual子結構，因此layers這個列表就是前面resnet50的[3, 4, 6, 3]。

　　　　- _make_layer方法中比較重要的兩行代碼是：

　　　　　　1、layers.append(block(self.inplanes, planes, stride, downsample))，

　　　　　　　    - 該部分是將每個blocks的第一個residual結構保存在layers列表中。

　　　　　　2、 for i in range(1, blocks): layers.append(block(self.inplanes, planes))，

　　　　　　　　- 該部分是將每個blocks的剩下residual 結構保存在layers列表中，這樣就完成了一個blocks的構造。

　　　　　　- 這兩行代碼中都是通過Bottleneck這個類來完成每個residual的構建。

2）接下來介紹Bottleneck類。

從前面的ResNet類可以看出，

在構造ResNet網絡的時候，

最重要的是Bottleneck這個類，

因為ResNet是由residual結構組成的，

而Bottleneck類就是完成residual結構的構建。

Bottlenect繼承了torch.nn.Module類，重寫了__init__()和forward()。

從forward方法可以看出，bottleneck就是我們熟悉的3個主要的卷積層、BN層和激活層，

最后的out += residual就是element-wise add的操作。

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

 

3）BasicBlock類和Bottleneck類類似，

前者主要是用來構建ResNet18和ResNet34網絡，

因為這兩個網絡的residual結構只包含兩個卷積層，沒有Bottleneck類中的bottleneck概念。

因此在該類中，第一個卷積層采用的是kernel_size=3的卷積，如conv3x3函數所示。

def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

 

4. 如何獲取預訓練模型。

前面提到這一行代碼：if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))，

主要就是通過model_zoo.py中的load_url函數根據model_urls字典導入相應的預訓練模型，models_zoo.py腳本的github地址：https://github.com/pytorch/pytorch/blob/master/torch/utils/model_zoo.py。

load_url函數源碼如下。

def load_url(url, model_dir=None, map_location=None, progress=True):
    r"""Loads the Torch serialized object at the given URL.

    If the object is already present in `model_dir`, it's deserialized and
    returned. The filename part of the URL should follow the naming convention
    ``filename-<sha256>.ext`` where ``<sha256>`` is the first eight or more
    digits of the SHA256 hash of the contents of the file. The hash is used to
    ensure unique names and to verify the contents of the file.

    The default value of `model_dir` is ``$TORCH_HOME/models`` where
    ``$TORCH_HOME`` defaults to ``~/.torch``. The default directory can be
    overriden with the ``$TORCH_MODEL_ZOO`` environment variable.

    Args:
        url (string): URL of the object to download
        model_dir (string, optional): directory in which to save the object
        map_location (optional): a function or a dict specifying how to remap storage locations (see torch.load)
        progress (bool, optional): whether or not to display a progress bar to stderr

    Example:
        >>> state_dict = torch.utils.model_zoo.load_url('https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth')

    """
    if model_dir is None:
        torch_home = os.path.expanduser(os.getenv('TORCH_HOME', '~/.torch'))
        model_dir = os.getenv('TORCH_MODEL_ZOO', os.path.join(torch_home, 'models'))
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
    parts = urlparse(url)
    filename = os.path.basename(parts.path)
    cached_file = os.path.join(model_dir, filename)
    if not os.path.exists(cached_file):
        sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
        hash_prefix = HASH_REGEX.search(filename).group(1)
        _download_url_to_file(url, cached_file, hash_prefix, progress=progress)
    return torch.load(cached_file, map_location=map_location)

 

首先model_dir是下載下來的模型的保存地址，如果沒有指定的話就會保存在項目的.torch目錄下，最好指定。

cached_file是保存模型的路徑加上模型名稱。

接下來的 if not os.path.exists(cached_file)語句用來判斷是否指定目錄下已經存在要下載模型，

如果已經存在，就直接調用torch.load接口導入模型，

如果不存在，則從網上下載，下載是通過_download_url_to_file(url, cached_file, hash_prefix, progress=progress)進行的。

重點在於模型導入是通過torch.load()接口來進行的，不管你的模型是從網上下載的還是本地已有的。