【学习笔记】resnet-18 pytorch源代码解读

使用的resnet-18的源代码来源于 PyTorch1.0， torchvision0.2.2

ResNet-18网络结构简图

ResNet（Residual Neural Network）来源于微软研究院的Kaiming He等人的论文《Deep Residual Learning for Image Recognition》。ResNet-18的网络简图如下图（假设网络的输入的张量的形状为$3\times 64\times 64$）

如图resnet的结构分为四个stage，完整的ResNet-18的结构图在最后。

ResNet-18的代码结构

pytorch中定义了resnet-18，resnet-34，resnet-50，resnet-101，resnet-152，在pytorch中使用resnet-18的方法如下：

from torchvision import models
resnet = models.resnet18(pretrained=True)

其中pretrained参数表示是否载入在ImageNet上预训练的模型。通过models.resnet18函数载入网络模型，该函数的定义如下

def resnet18(pretrained=False, **kwargs):
    """构建一个ResNet-18模型

    参数:
        pretrained (bool): 若为True则返回在ImageNet上预训练的模型
    """
    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
    return model

使用其他的ResNet模型时，使用对应的名字的函数就行，函数返回的是一个ResNet类型的实例，这个类定义了resnet网络的结构，ResNet类的结构如下：

class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False):
        """定义ResNet网络的结构
        
        参数：
            block (BasicBlock / Bottleneck): 残差块类型
            layers (list): 每一个stage的残差块的数目，长度为4
            num_classes (int): 类别数目
            zero_init_residual (bool): 若为True则将每个残差块的最后一个BN层初始化为零，
                这样残差分支从零开始每一个残差分支，每一个残差块表现的就像一个恒等映射，根据
                https://arxiv.org/abs/1706.02677这可以将模型的性能提升0.2~0.3%
        """
        super(ResNet, self).__init__()
        # __init__
    def _make_layer(self, block, planes, blocks, stride=1):
        # _make_layer function
    def forward(self, x):
        # forward function

残差块结构

注意上面的ResNet类的代码中的block参数，这个参数定义了残差块的结构，分为两种：

1. BasicBlock：resnet-18和resnet-34的残差块结构；
2. Bottleneck：resnet-50，resnet-101和resnet-152的残差块结构；

这里仅关注resnet-18因此我们只关注BasicBlock类，这个类的结构如下：

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)


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        """定义BasicBlock残差块类
        
        参数：
            inplanes (int): 输入的Feature Map的通道数
            planes (int): 第一个卷积层输出的Feature Map的通道数
            stride (int, optional): 第一个卷积层的步长
            downsample (nn.Sequential, optional): 旁路下采样的操作
        注意：
            残差块输出的Feature Map的通道数是planes*expansion
        """
        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):
        identity = 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:
            identity = self.downsample(x)

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

        return out

首先BasicBlock类有一个类属性，BasicBlock.expansion这个类属性的值为1，另外在 Bottleneck类中也有这个类属性，值为4，这个类属性表示残差块的输入的Feature Map的通道数为$inplanes$，输出的通道数为$planes\times expansion$
$${\text{out}}_{(planes\times expansion)\times H^{\prime }\times W^{\prime }}=Block({\text{x}}_{inplanes\times H\times W})$$

BasicBlock定义的网络结构如下所示：

ResNet类

构造方法和forward

首先是构造方法和forward方法

class ResNet(nn.Module):

    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False):
         """定义ResNet网络的结构
        
        参数：
            block (BasicBlock / Bottleneck): 残差块类型
            layers (list): 每一个stage的残差块的数目，长度为4
            num_classes (int): 类别数目
            zero_init_residual (bool): 若为True则将每个残差块的最后一个BN层初始化为零，
                这样残差分支从零开始每一个残差分支，每一个残差块表现的就像一个恒等映射，根据
                https://arxiv.org/abs/1706.02677这可以将模型的性能提升0.2~0.3%
        """
        super(ResNet, self).__init__()
        self.inplanes = 64  # 第一个残差块的输入通道数
        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)
        
        # Stage1 ~ Stage4
        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.AdaptiveAvgPool2d((1, 1))  # GAP
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        # 网络参数初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)
    
    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

ResNet类定义的网络结构如下图所示：

_make_layer方法

ResNet类中还有一个方法是_make_layer在这个方法中，定义了ResNet网络的一个Stage的结构，代码如下：

def _make_layer(self, block, planes, blocks, stride=1):
    """定义ResNet的一个Stage的结构
    
    参数：
        block (BasicBlock / Bottleneck): 残差块结构
        plane (int): 残差块中第一个卷积层的输出通道数
        bloacks (int): 当前Stage中的残差块的数目
        stride (int): 残差块中第一个卷积层的步长
    """
    downsample = None
    if stride != 1 or self.inplanes != planes * block.expansion:
        downsample = nn.Sequential(
            conv1x1(self.inplanes, planes * block.expansion, stride),
            nn.BatchNorm2d(planes * block.expansion),
        )

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

    return nn.Sequential(*layers)

首先是为第一个残差块定义downsample结构，当残差块的输入和输出的尺寸不一致或者通道数不一致的时候就会需要下采样结构，下采样结构由一个1x1卷积层和一个BatchNorm层组成。之后定义了$blocks$个残差块（在ResNet-18中每一个Stage均有两个残差块），只有第一个残差块需要下采样层。

在ResNet中，Stage1中输入通道数和输出通道数相同，并且使用的是stride为1的卷积，因此在Stage1中不需要有下采样层，其余Stage中钧需要有下采样层。

完整的ResNet-18结构图

以输入$3\times 64\times 64$图像为例。