【Faster R-CNN】之 backbone 代码精读

1、前言

在上一篇文章 【Faster R-CNN】之 Resize_and_Padding 代码精读 中，我们得到了图像尺寸统一的 batch 了，接下来，就是 feed 到 backbone 中 获取 feature map。

2、backbone

1）这里的 backbone 其实就是 迁移学习，通过加载已有的一些 经典模型 和 其在大数据集上训练好的参数，来提取我们当前图像中的特征 (feature map)。

2）我们不会使用 加载下来的 经典模型的全部，而是使用模型其中的一部分（前部），主要是不会使用 模型的 全连接层 那一部分，因为 我们想要的 图像特征需要是二维的形式，全连接层 是一维的，我不能用。
所以，我们在 load 模型之后，会只截取 模型的前部分 用于 图像特征提取

3）我们这里为了讲解简单，采用不带 fpn 的形式。等到不带 fpn 的 faster r-cnn 学习完，再加上 fpn 看，也就简单多了。

4）加载模型时，torchvision.models 函数有 新写法 和老写法，老写法使用的参数形式为 pretrained=True ， 可能会 warning。如果你 torchvision 版本在 0.13之后，建议就按照 新写法使用了。 具体参考文章 看这里

---

这里我们以 resnet18 作为 backbone 来做代码讲解。
** 为啥我用 resent18呢？ 因为 resnet50 模型太大了，我的 gpu 承载不了。。。 而且 resnet 的效果也不错啦，大家也可以试试其他的， mobilenet，vgg16， efficientNet 等

代码讲解

import torchvision.models as models
backbone = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

• model = torchvision.models.resent18() 用于加载模型，
• 参数 weights=models.ResNet18_Weights.DEFAULT 表示加载模型的同时，一并加载 训练好的模型参数，模型参数的 DEFAULT 版本，一般都是效果最好的最新版本

把 加载的 resnet 模型 print 出来看下

ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (1): BasicBlock(
      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer2): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer3): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer4): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): Linear(in_features=512, out_features=1000, bias=True)
)

我们用 layer4 的输出结果，作为图像的特征

from torchvision.models.feature_extraction import create_feature_extractor
backbone = create_feature_extractor(backbone, return_nodes={"layer4": "0"})

我们使用 torchvision.models.feature_extraction(backbone, return_nodes={"layer4": "0"}) 函数 截取 resnet18 前半部分 到 layer4 的模型。

我们随机生成一张图片，放到 backbone中，看下输出的结果是什么样

out = backbone(torch.rand(1, 3, 224, 224))
print(out)

输出是一个 字典，key 为我们指定的 ‘0’，value 就是 我们要的 feature了，shape 为 (1， 512， 7， 7)。
输出的feature 的 channel 数为 512， feature map 的大小为 7x7

到这里，我们以 resnet18 为基础创建 backbone 就搞定了。
最后再指定下 out_channels 属性，方便后面使用。

backbone.out_channels = 512

---

3、代码

import torch
import torchvision.models as models
from torchvision.models.feature_extraction import create_feature_extractor


def resnet18_backbone():
    backbone = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
    # print(backbone)
    backbone = create_feature_extractor(backbone, return_nodes={"layer4": "0"})
    out = backbone(torch.rand(1, 3, 224, 224))
    # print(out["0"].shape)
    backbone.out_channels = 512
    return backbone


def vgg16_backbone():
    backbone = models.vgg16_bn(pretrained=True)
    # print(backbone)
    backbone = create_feature_extractor(backbone, return_nodes={"features.42": "0"})
    # out = backbone(torch.rand(1, 3, 224, 224))
    # print(out["0"].shape)
    backbone.out_channels = 512
    return backbone


def efficientNetB0_backbone():
    backbone = models.efficientnet_b0(pretrained=True)
    # print(backbone)
    backbone = create_feature_extractor(backbone, return_nodes={"features.5": "0"})
    # out = backbone(torch.rand(1, 3, 224, 224))
    # print(out["0"].shape)
    backbone.out_channels = 112
    return backbone