ResNet程式碼精讀

只有菜的小白發表於2024-05-02
class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_channel, out_channel, stride=1, downsample=None, **kwargs):  # 虛線對應的 downsample
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,
                               kernel_size=3, stride=stride, padding=1, bias=False) # 有BN層不需要偏置 ,這裡的stride需要根據傳進來的值對矩陣改變寬高
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel, # 第一個卷積已經改變寬高
                               kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)
        self.downsample = downsample

    def forward(self, x):
        identity = x  # 分支線上的
        if self.downsample is not None:
            identity = self.downsample(x)

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

        out = self.conv2(out)
        out = self.bn2(out) # 先加上捷徑分支,再relu啟用

        out += identity  # 加上捷徑的再輸出
        out = self.relu(out)

        return out

上面定義了一個兩層的卷積層,論文中有用到過。

class Bottleneck(nn.Module):
    """
    注意:原論文中,在虛線殘差結構的主分支上,第一個1x1卷積層的步距是2,第二個3x3卷積層步距是1。
    但在pytorch官方實現過程中是第一個1x1卷積層的步距是1,第二個3x3卷積層步距是2,
    這麼做的好處是能夠在top1上提升大概0.5%的準確率。
    可參考Resnet v1.5 https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch
    """
    expansion = 4

    def __init__(self, in_channel, out_channel, stride=1, downsample=None,
                 groups=1, width_per_group=64):
        super(Bottleneck, self).__init__()

        width = int(out_channel * (width_per_group / 64.)) * groups

        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=width,           # 這裡的out_channel是第一層和第二層的輸出矩陣深度
                               kernel_size=1, stride=1, bias=False)  # squeeze channels # 第二層的stride根據傳入的來判斷
        self.bn1 = nn.BatchNorm2d(width)
        # -----------------------------------------
        self.conv2 = nn.Conv2d(in_channels=width, out_channels=width, groups=groups,
                               kernel_size=3, stride=stride, bias=False, padding=1) # stride是為了調整矩陣的寬高
        self.bn2 = nn.BatchNorm2d(width)
        # -----------------------------------------
        self.conv3 = nn.Conv2d(in_channels=width, out_channels=out_channel*self.expansion,
                               kernel_size=1, stride=1, bias=False)  # unsqueeze channels
        self.bn3 = nn.BatchNorm2d(out_channel*self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        identity = x
        if self.downsample is not None:
            identity = self.downsample(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)

        out += identity  #先加上再啟用
        out = self.relu(out)

        return out


這個殘差塊和第一個殘差塊我覺得唯一的區別就是這個殘差塊多了一層。
這裡提一下一直以來都比較困惑的一個點:虛線殘差塊有兩個作用,一個是改變矩陣深度,如resnet50,101,152的conv_2x的第一層,都是改變輸入矩陣的深度,而不改變輸入矩陣的寬高。因為輸入輸出矩陣的寬高是一致的。而且虛線殘差塊僅出現在每個conv_··x的第一層,因為經過第一層之後,矩陣的深度和寬高都被調整為對應的輸出矩陣的寬高,所以後面的都是實線殘差結構。這也就是為什麼下面的for迴圈可以直接將剩下的殘差塊壓入。

class ResNet(nn.Module):

    def __init__(self,
                 block,   # 根據模型選擇bottleneck還是basicmodule
                 blocks_num,
                 num_classes=1000,
                 include_top=True,
                 groups=1,
                 width_per_group=64):
        super(ResNet, self).__init__()
        self.include_top = include_top
        self.in_channel = 64

        self.groups = groups
        self.width_per_group = width_per_group

        self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2,
                               padding=3, bias=False)  # 設定padding是為了高和寬縮減為原來的一半
        self.bn1 = nn.BatchNorm2d(self.in_channel)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)  # 設定padding理由同上
        self.layer1 = self._make_layer(block, 64, blocks_num[0])
        self.layer2 = self._make_layer(block, 128, blocks_num[1], stride=2)
        self.layer3 = self._make_layer(block, 256, blocks_num[2], stride=2)
        self.layer4 = self._make_layer(block, 512, blocks_num[3], stride=2)
        if self.include_top:
            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # output size = (1, 1) 自適應展平為長寬(1*1)的矩陣
            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')

    def _make_layer(self, block, channel, block_num, stride=1): # channel都是第一層的深度,但是50層和以上都有4倍的最後一層
        downsample = None
        if stride != 1 or self.in_channel != channel * block.expansion: # 第二個條件是因為resnet50,101,152的conv_2x雖然stride是1(輸入輸出矩陣同寬高),但是深度不同,所以第一層也得 
                                                                        #  加上虛線殘差結構
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False), # 對於layer1來說,因為池化過後的長寬和輸出的長寬一致,所以也就虛線stride就是1,而對於下面的三層,stride都是2才能改變矩陣的長寬
                nn.BatchNorm2d(channel * block.expansion))

        layers = []
        layers.append(block(self.in_channel,
                            channel,
                            downsample=downsample,
                            stride=stride,
                            groups=self.groups,
                            width_per_group=self.width_per_group))
        self.in_channel = channel * block.expansion    #這裡的in_channel是個成員資料,對於每個殘差塊,輸出矩陣深度不一致。對於下面的迴圈,輸入輸出矩陣深度應該一致(實線殘差結構)。

        for _ in range(1, block_num):  # 這裡剩下的都是實線殘差結構了,因此直接加上即可
            layers.append(block(self.in_channel,
                                channel,
                                groups=self.groups,
                                width_per_group=self.width_per_group))

        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)

        if self.include_top:
            x = self.avgpool(x)
            x = torch.flatten(x, 1)
            x = self.fc(x)

        return x

給個簡單的模型構建傳參的例子

def resnet50(num_classes=1000, include_top=True):
    # https://download.pytorch.org/models/resnet50-19c8e357.pth
    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top)


根據block判斷選擇幾層的殘差結構

相關文章