Faster RCNN中Backbone（ResNet50+FPN）

Backbone

这部分展示了使用ResNet50+FPN作为Faster RCNN的特征提取部分

import os
from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F


class IntermediateLayerGetter(nn.ModuleDict):
    def __init__(self, model, return_layers):
        # 首先判断 return_layer中的key 是否在model中
        if not set(return_layers).issubset([name for name, _ in model.named_children()]):
            raise ValueError("return_layers are not present in model")

        orig_return_layers = return_layers
        return_layers = {str(k): str(v) for k, v in return_layers.items()}
        layers = OrderedDict()

        # 遍历模型子模块按顺序存入有序字典
        # 只保存layer4及其之前的结构，舍去之后不用的结构
        for name, module in model.named_children():
            layers[name] = module
            if name in return_layers:
                del return_layers[name]
            if not return_layers:
                break

        super(IntermediateLayerGetter, self).__init__(layers)
        self.return_layers = orig_return_layers

    def forward(self, x):
        out = OrderedDict()
        # 依次遍历模型的所有子模块，并进行正向传播，
        # 收集layer1, layer2, layer3, layer4的输出
        for name, module in self.items():
            x = module(x)
            if name in self.return_layers:
                out_name = self.return_layers[name]
                out[out_name] = x
        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, in_channel, out_channel, stride=1, downsample=None, norm_layer=nn.BatchNorm2d):
        super(Bottleneck, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=1, stride=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)

        self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)

        self.conv3 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel * self.expansion, kernel_size=1, stride=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channel * self.expansion)

        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

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

        return out


class ResNet(nn.Module):
    # 仅包含分类器前面的模型结构
    def __init__(self, block, blocks_num, num_classes=1000):
        super(ResNet, self).__init__()

        self._norm_layer = nn.BatchNorm2d

        self.in_channel = 64

        self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(self.in_channel)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        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)

        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):
        norm_layer = self._norm_layer
        downsample = None
        if stride != 1 or self.in_channel != channel * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),
                norm_layer(channel * block.expansion))

        layers = []
        layers.append(block(self.in_channel, channel, downsample=downsample,
                            stride=stride, norm_layer=norm_layer))
        self.in_channel = channel * block.expansion

        for _ in range(1, block_num):
            layers.append(block(self.in_channel, channel, norm_layer=norm_layer))

        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)

        return x


class MaxPool_Block(nn.Module):

    def forward(self, x, names):
        names.append('pool')
        # input, kernel_size, stride, padding
        x.append(F.max_pool2d(x[-1], 1, 2, 0))
        return x, names


class FPN(nn.Module):
    def __init__(self, backbone_output_channels_list, out_channels, maxpool_blocks=True):
        super(FPN, self).__init__()

        # 用来调整resnet输出特征矩阵(layer1,2,3,4)的channel（kernel_size=1）
        self.inner_blocks = nn.ModuleList()
        # 对调整后的特征矩阵使用3x3的卷积核来得到对应的预测特征矩阵
        self.layer_blocks = nn.ModuleList()

        for backbone_channels in backbone_output_channels_list:
            inner_block_module = nn.Conv2d(backbone_channels, out_channels, 1)
            layer_block_module = nn.Conv2d(out_channels, out_channels, 3, padding=1)

            self.inner_blocks.append(inner_block_module)
            self.layer_blocks.append(layer_block_module)
        # 初始化权重参数
        for m in self.children():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_uniform_(m.weight, a=1)
                nn.init.constant_(m.bias, 0)

        self.extra_blocks = maxpool_blocks

    def forward(self, x):
        # 对应backbone中layer1、2、3、4特征图的输出的key
        names = list(x.keys())
        # 对应backbone中layer1、2、3、4特征图的输出的value
        x = list(x.values())
        # backbone中layer4的输出调整到指定维度上 channels: 2048->256
        last_inner = self.inner_blocks[-1](x[-1])
        # results中保存着经过3*3conv后的每个预测特征层
        results = []
        # layer4
        results.append(self.layer_blocks[-1](last_inner))
        # layer4+layer3
        layer3_inner = self.inner_blocks[2](x[2])
        layer3_shape = layer3_inner.shape[-2:]
        last_layer3 = F.interpolate(last_inner, size=layer3_shape, mode="nearest")
        last_layer3_sum = layer3_inner + last_layer3
        results.insert(0, self.layer_blocks[2](last_layer3_sum))
        # layer3+layer2
        layer2_inner = self.inner_blocks[1](x[1])
        layer2_shape = layer2_inner.shape[-2:]
        layer3_layer2 = F.interpolate(last_layer3_sum, size=layer2_shape, mode="nearest")
        layer3_layer2_sum = layer2_inner + layer3_layer2
        results.insert(0, self.layer_blocks[1](layer3_layer2_sum))
        # layer2+layer1
        layer1_inner = self.inner_blocks[0](x[0])
        layer1_shape = layer1_inner.shape[-2:]
        layer2_layer1 = F.interpolate(layer3_layer2_sum, size=layer1_shape, mode="nearest")
        layer2_layer1_sum = layer1_inner + layer2_layer1
        results.insert(0, self.layer_blocks[0](layer2_layer1_sum))
        # results 存储着FPN后特征图从大到小的 key:0,1,2,3,4. value(H W shape):[1/4, 1/8, 1/16, 1/32, 1/64]
        if self.extra_blocks:
            results, names = self.extra_blocks(results, names)
        # out: key: 0,1,2,3,pool
        out = OrderedDict([(k, v) for k, v in zip(names, results)])

        return out


class Backbone_FPN(nn.Module):
    def __init__(self, backbone, return_layers, in_channels_list, out_channels):
        super(Backbone_FPN, self).__init__()

        maxpool_block = MaxPool_Block()
        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
        self.fpn = FPN(
            backbone_output_channels_list=in_channels_list,
            out_channels=out_channels,
            maxpool_blocks=maxpool_block,
        )

        self.out_channels = out_channels

    def forward(self, x):
        x = self.body(x)
        x = self.fpn(x)

        return x


def ResNet50_FPN_backbone():
    resnet_backbone = ResNet(Bottleneck, [3, 4, 6, 3])
    return_layers = {'layer1': '0', 'layer2': '1', 'layer3': '2', 'layer4': '3'}
    in_channels_list = [256, 512, 1024, 2048]
    out_channels = 256

    return Backbone_FPN(resnet_backbone, return_layers, in_channels_list, out_channels)


if __name__ == '__main__':
    model = ResNet50_FPN_backbone()
    print(model)