Pytorch 对 ResNet 的实现

参考

借助 torch hub

可以直接借助 torch hub pytorch vision resnet 来加载预训练的 resnet。

import torch
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
'''or any of these variants'''
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet34', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet101', pretrained=True)
# model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet152', pretrained=True)
model.eval()

这种直接调用 torch hub 接口的方式，优点是很简单易用，缺点是它默认的输入是 [B, 3, H, W]，默认的输出是 [B, 1000]，不提供修改 in_channels 和 num_classes 的接口。

重写 torch resnet

参考 Pytorch 对 resnet 的实现代码 resnet.py | pytorch github，对代码进行简化。这种 ResNet 实现的主体结构与 Pytorch 的一致，可以直接导入 Pytorch 官方提供的 ResNet 预训练模型。

'''ref: https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py'''

from typing import Any, Callable, List, Optional, Type, Union

import torch
import torch.nn as nn
from torch import Tensor

RESNET_WEIGHTS_URL = {
    'IMAGENET1k_V1': {
        'resnet18': "https://download.pytorch.org/models/resnet18-f37072fd.pth",
        'resnet34': "https://download.pytorch.org/models/resnet34-b627a593.pth",
        'resnet50': "https://download.pytorch.org/models/resnet50-0676ba61.pth",
        'resnet101': "https://download.pytorch.org/models/resnet101-63fe2227.pth",
        'resnet152': "https://download.pytorch.org/models/resnet152-394f9c45.pth",
    },

    'IMAGENET1k_V2': {
        'resnet50': "https://download.pytorch.org/models/resnet50-11ad3fa6.pth",
        'resnet101': "https://download.pytorch.org/models/resnet101-cd907fc2.pth",
        'resnet152': "https://download.pytorch.org/models/resnet152-f82ba261.pth",
    },
}

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


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


class BasicBlock(nn.Module):
    expansion: int = 1

    def __init__(
        self,
        inplanes: int,
        planes: int,
        stride: int = 1,
        downsample: Optional[nn.Module] = None,
        groups: int = 1,
        base_width: int = 64,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None,
    ) -> None:
        super().__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if groups != 1 or base_width != 64:
            raise ValueError("BasicBlock only supports groups=1 and base_width=64")
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        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


class Bottleneck(nn.Module):
    """
    Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
    while original implementation places the stride at the first 1x1 convolution(self.conv1)
    according to "Deep residual learning for image recognition" https://arxiv.org/abs/1512.03385.
    This variant is also known as ResNet V1.5 and improves accuracy according to
    https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
    """

    expansion: int = 4

    def __init__(
        self,
        inplanes: int,
        planes: int,
        stride: int = 1,
        downsample: Optional[nn.Module] = None,
        groups: int = 1,
        base_width: int = 64,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None,
    ) -> None:
        super().__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.0)) * groups
        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

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

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

        return out


class ResNet(nn.Module):
    def __init__(
        self,
        block: Type[Union[BasicBlock, Bottleneck]],
        layers: List[int],
        num_classes: int = 1000,
        zero_init_residual: bool = False,
        groups: int = 1,
        width_per_group: int = 64,
        replace_stride_with_dilation: Optional[List[bool]] = None,
        norm_layer: Optional[Callable[..., nn.Module]] = None,
    ) -> None:
        super().__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = 64
        self.dilation = 1
        if replace_stride_with_dilation is None:
            # each element in the tuple indicates if we should replace
            # the 2x2 stride with a dilated convolution instead
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError(
                "replace_stride_with_dilation should be None "
                f"or a 3-element tuple, got {replace_stride_with_dilation}"
            )
        self.groups = groups
        self.base_width = width_per_group
        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = norm_layer(self.inplanes)
        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, dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1])
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2])
        self.avgpool = nn.AdaptiveAvgPool2d((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")
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck) and m.bn3.weight is not None:
                    nn.init.constant_(m.bn3.weight, 0)  # type: ignore[arg-type]
                elif isinstance(m, BasicBlock) and m.bn2.weight is not None:
                    nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]

    def _make_layer(
        self,
        block: Type[Union[BasicBlock, Bottleneck]],
        planes: int,
        blocks: int,
        stride: int = 1,
        dilate: bool = False,
    ) -> nn.Sequential:
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                norm_layer(planes * block.expansion),
            )

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

        return nn.Sequential(*layers)

    def _forward_impl(self, x: Tensor) -> Tensor:
        # See note [TorchScript super()]
        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 = torch.flatten(x, 1)
        x = self.fc(x)

        return x

    def forward(self, x: Tensor) -> Tensor:
        return self._forward_impl(x)


def _resnet(
    block: Type[Union[BasicBlock, Bottleneck]],
    layers: List[int],
    num_classes: int,
    **kwargs: Any,
) -> ResNet:
    model = ResNet(block, layers, num_classes=num_classes, **kwargs)
    return model


def resnet18(*, num_classes: int = 1000, **kwargs: Any) -> ResNet:
    """ResNet-18 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`.

    Args:
        num_classes (int, optional): number of classes. Default is 1000.
        **kwargs: parameters passed to the `ResNet` base class.
    """
    # weights = ResNet18_Weights.verify(weights)

    return _resnet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes, **kwargs)


def resnet34(*, num_classes: int = 1000, **kwargs: Any) -> ResNet:
    """ResNet-34 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`.

    Args:
        num_classes (int, optional): number of classes. Default is 1000.
        **kwargs: parameters passed to the `ResNet` base class.
    """

    return _resnet(BasicBlock, [3, 4, 6, 3], num_classes=num_classes, **kwargs)


def resnet50(*, num_classes: int = 1000, **kwargs: Any) -> ResNet:
    """ResNet-50 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`.

    Note:
       The bottleneck of TorchVision places the stride for downsampling to the second 3x3
       convolution while the original paper places it to the first 1x1 convolution.
       This variant improves the accuracy and is known as `ResNet V1.5
       <https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch>`.

    Args:
        num_classes (int, optional): number of classes. Default is 1000.
        **kwargs: parameters passed to the `ResNet` base class.
    """

    return _resnet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, **kwargs)


def resnet101(*, num_classes: int = 1000, **kwargs: Any) -> ResNet:
    """ResNet-101 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`.

    Note:
       The bottleneck of TorchVision places the stride for downsampling to the second 3x3
       convolution while the original paper places it to the first 1x1 convolution.
       This variant improves the accuracy and is known as `ResNet V1.5
       <https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch>`.

    Args:
        num_classes (int, optional): number of classes. Default is 1000.
        **kwargs: parameters passed to the `ResNet` base class.
    """

    return _resnet(Bottleneck, [3, 4, 23, 3], num_classes=num_classes, **kwargs)


def resnet152(*, num_classes: int = 1000, **kwargs: Any) -> ResNet:
    """ResNet-152 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`.

    Note:
        The bottleneck of TorchVision places the stride for downsampling to the second 3x3
        convolution while the original paper places it to the first 1x1 convolution.
        This variant improves the accuracy and is known as `ResNet V1.5
        <https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch>`.

    Args:
        num_classes (int, optional): number of classes. Default is 1000.
        **kwargs: parameters passed to the `ResNet` base class.
    """

    return _resnet(Bottleneck, [3, 8, 36, 3], num_classes=num_classes, **kwargs)

if __name__ == "__main__":
    import torch.utils.model_zoo as model_zoo

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    B, C, H, W = 8, 3, 224, 224
    input = torch.randn(size=[B, C, H, W], device=device)

    model = resnet18().to(device)
    # model = resnet34().to(device)
    # model = resnet50().to(device)
    # model = resnet101().to(device)
    # model = resnet152().to(device)
    print(model)

    weights_url = RESNET_WEIGHTS_URL["IMAGENET1k_V1"]["resnet18"]
    pretrained_dict = model_zoo.load_url(url=weights_url)
    model.load_state_dict(pretrained_dict)
    
    output = model(input)
    print(output.shape)

预训练权重

使用 torch.utils.model_zoo 来从对应的 url 中下载和加载预训练权重，下面是对 resnet.py | pytorch github 中预训练权重的整理

RESNET_WEIGHTS_URL = {
  'IMAGENET1k_V1': {
    'resnet18': "https://download.pytorch.org/models/resnet18-f37072fd.pth",
    'resnet34': "https://download.pytorch.org/models/resnet34-b627a593.pth",
    'resnet50': "https://download.pytorch.org/models/resnet50-0676ba61.pth",
    'resnet101': "https://download.pytorch.org/models/resnet101-63fe2227.pth",
    'resnet152': "https://download.pytorch.org/models/resnet152-394f9c45.pth",
  },

  'IMAGENET1k_V2': {
    'resnet50': "https://download.pytorch.org/models/resnet50-11ad3fa6.pth",
    'resnet101': "https://download.pytorch.org/models/resnet101-cd907fc2.pth",
    'resnet152': "https://download.pytorch.org/models/resnet152-f82ba261.pth",
  },
}

提取 features

很多模型中会使用 ResNet 提取特征，我们还可以在 ResNet 实现的 forward 中添加获取中间 features 的代码，并将结果进行输出

class ResNet(nn.Module):
    ...
    def _forward_impl(self, x: Tensor) -> Tuple[Tensor, List[Tensor]]:
        # See note [TorchScript super()]
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        feature0 = x

        x = self.layer1(x)
        feature1 = x
        x = self.layer2(x)
        feature2 = x
        x = self.layer3(x)
        feature3 = x
        x = self.layer4(x)
        feature4 = x

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

        return x, [feature0, feature1, feature2, feature3, feature4]

    def forward(self, x: Tensor) -> Tuple[Tensor, List[Tensor]]:
        return self._forward_impl(x)

if __name__ == "__main__":
    import torch.utils.model_zoo as model_zoo

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    B, C, H, W = 8, 3, 224, 224
    input = torch.randn(size=[B, C, H, W], device=device)

    model = resnet18().to(device)
    # model = resnet34().to(device)
    # model = resnet50().to(device)
    # model = resnet101().to(device)
    # model = resnet152().to(device)
    print(model)

    weights_url = RESNET_WEIGHTS_URL["IMAGENET1k_V1"]["resnet18"]
    pretrained_dict = model_zoo.load_url(url=weights_url)
    model.load_state_dict(pretrained_dict)
    
    output, features = model(input)
    print(output.shape)
    print(f"{features[0].shape = }\n" # [B,  64, 56, 56]
          f"{features[1].shape = }\n" # [B,  64, 56, 56]
          f"{features[2].shape = }\n" # [B, 128, 28, 28]
          f"{features[3].shape = }\n" # [B, 256, 14, 14]
          f"{features[4].shape = }")  # [B, 512,  7,  7]

使用 IntermediateLayerGetter 提取 features

此外还可以借助 IntermediateLayerGetter 方法来提取中间结果。关于 IntermediateLayerGetter 方法更详细的介绍请见 IntermediateLayerGetter 获取中间结果 | 文羊羽，此处仅做展示：

import torch
from torchvision.models import resnet18, ResNet18_Weights
from torchvision.models._utils import IntermediateLayerGetter
input = torch.rand(size=(1, 3, 224, 224))

model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
output = model(input)
print(output.shape)
# torch.Size([1, 1000])

'''extract layer1, layer2, layer3 and layer4, giving as names feat1, feat2, feat3, feat4 respectively'''
return_layers = {'layer1':'feat1', 'layer2':'feat2', 'layer3':'feat3', 'layer4':'feat4'}
new_model = IntermediateLayerGetter(model, return_layers)
output = new_model(input)
for k, v in output.items():
  print(k, v.shape)
# feat1 torch.Size([1, 64, 56, 56])
# feat2 torch.Size([1, 128, 28, 28])
# feat3 torch.Size([1, 256, 14, 14])
# feat4 torch.Size([1, 512, 7, 7])

ResNet 所提取 features 的 Shape

resnet18 / resnet34

input: [B, 3, 224, 224]

features[0].shape = torch.Size([B,  64, 56, 56])
features[1].shape = torch.Size([B,  64, 56, 56])
features[2].shape = torch.Size([B, 128, 28, 28])
features[3].shape = torch.Size([B, 256, 14, 14])
features[4].shape = torch.Size([B, 512,  7,  7])

input: [B, 3, 256, 256]

features[0].shape = torch.Size([B,  64, 64, 64])
features[1].shape = torch.Size([B,  64, 64, 64])
features[2].shape = torch.Size([B, 128, 32, 32])
features[3].shape = torch.Size([B, 256, 16, 16])
features[4].shape = torch.Size([B, 512,  8,  8])

input: [B, 3, 512, 512]

features[0].shape = torch.Size([B,  64, 128, 128])
features[1].shape = torch.Size([B,  64, 128, 128])
features[2].shape = torch.Size([B, 128,  64,  64])
features[3].shape = torch.Size([B, 256,  32,  32])
features[4].shape = torch.Size([B, 512,  16,  16])

input: [B, 3, 1024, 1024]

features[0].shape = torch.Size([B,  64, 256, 256])
features[1].shape = torch.Size([B,  64, 256, 256])
features[2].shape = torch.Size([B, 128, 128, 128])
features[3].shape = torch.Size([B, 256,  64,  64])
features[4].shape = torch.Size([B, 512,  32,  32])

input: [B, 3, 2048, 2048]

features[0].shape = torch.Size([B,  64, 512, 512])
features[1].shape = torch.Size([B,  64, 512, 512])
features[2].shape = torch.Size([B, 128, 256, 256])
features[3].shape = torch.Size([B, 256, 128, 128])
features[4].shape = torch.Size([B, 512,  64,  64])

resnet50 / resnet101 / resnet152