Faster RCNN代码整体框架

前言

Faster RCNN作为经典的双阶段目标检测算法，掌握其中的思想和代码实现的方法，对于我们实现单阶段目标检测或者双阶段目标检测都是很有帮助的。
相较于单阶段目标检测，双阶段目标检测主要多了一步生成proposal，也就是候选框的生成。在Faster RCNN中，对于图像中的生成的每一个anchor而言，首先要经过RPN（在这里只区分前景或者背景）做第一次筛选，选出概率大的一些anchor（train时采用预测目标概率前2000个anchor，test时采用1000个）作为proposal，proposal用于后续Faster RCNN的计算，这样的方式也就称为双阶段目标检测。

我以源码为基础，将代码更精简了一点，实现起来更加容易，并且细化了代码量并在每部分做了详细的注释以及shape的变化，以便更好的理解代码的原理和实现的过程。

整体组成

整体Faster RCNN主要包括transform、backbone、RegionProposalNetwork(RPN)、RoIHead、postprocess组成。
其中：transform的作用是将输入图像，缩放到一个给定的大小中，组成一个batch数据输入给网络。因为对于输入图像，其都是一张张大小不一尺寸的图像，无法输入到网络做平行计算，因此transform就是为了将输入图像缩放到指定大小，组成一个batch，然后输入到网络中。
backbone：对于输入图像进行特征提取，得到特征图。在这里采用的是ResNet50+FPN作为backbone。特征图包括ResNet50中layer1、layer2、layer3、layer4的输出并做FPN，将输出的通道数(256, 512, 1024, 2048)都调整为256，并且在layer4生成的特征图后使用一个maxpool，得到’pool’。一共生成5个预测特征图，分别为layer1, layer2, layer3, layer4, maxpool。对应的键为0, 1, 2, 3, pool.
RegionProposalNetwork: 将图像中生成的anchors进行筛选，生成proposal。
RoIHead: 由RoI pooling（RoIAlign pooling） + MLP(FastRCNN中ROI-pooling后的展平操作+两个全连接层) + FastRCNNPredictor(输出的预测部分，包括bbox回归参数和类别概率)组成。对RPN生成的proposal进行bbox回归参数和类别参数的预测。
postprocess：将网络的预测结果还原到原图像尺寸上。

Faster RCNN Model整体代码

import torch
from torch import nn, Tensor
import torch.nn.functional as F
from torchvision.ops import MultiScaleRoIAlign

from transforms import RcnnTransforms
from RPN import AnchorsGenerator, RPNHead, RegionProposalNetwork
from RoIHead import RoIHead


class FasterRCNNBase(nn.Module):

    def __init__(self, backbone, rpn, roi_heads, transform):
        super(FasterRCNNBase, self).__init__()
        self.transform = transform
        self.backbone = backbone
        self.rpn = rpn
        # roi pooling + MLP + FastRCNNPredictor + postprocess detections
        self.roi_heads = roi_heads

    def forward(self, images, targets=None):
        # 这里输入的images大小都是不同的，后面进行transform后修改成为同样大小的tensor打包成一个batch
        # images为dataloader的返回结果,list中每一个图像都是一个tensor类型
        # targets也是一个list，其里面每个元素是一个字典类型，包含每一个图像中的标注信息。
        original_image_sizes = []

        for img in images:
            val = img.shape[-2:] # B C H W
            assert len(val) == 2  # 防止输入的是个一维向量
            # 记录原始图像的size 因为transform后图像size发生变化，在最终的输出会映射回原图大小。
            original_image_sizes.append((val[0], val[1]))
        # 对图像进行预处理 得到新的 images targets。
        # 在输入transform之前，图像都是一张张的图像，图像大小不一。通过transform，将图像统一放到一个给定大小的tensor中，得到了一个batch的数据。
        images, targets = self.transform(images, targets)
        # 将图像输入backbone得到特征图
        features = self.backbone(images.tensors)

        # 将特征层以及标注target信息传入rpn中 -> 区域建议框proposal和RPN的损失loss
        # proposals shape: [num_proposals, 4] 4: xmin ymin xmax ymax
        proposals, proposal_losses = self.rpn(images, features, targets)

        # 将rpn生成的数据以及标注target信息传入fastRCNN后半部分
        detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)

        # 对网络的预测结果进行后处理（主要将bboxes还原到原图像尺度上）
        # 得到最终检测的一系列目标和FastRCNN的损失值
        # images.image_sizes 通过transform后的图像尺寸 original_image_sizes预处理前的图像尺寸。
        # 将预测结果还原到最初的图像尺寸上
        detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)

        losses = {}
        losses.update(detector_losses)
        losses.update(proposal_losses)

        if self.training:
            return losses

        return detections


class TwoMLPHead(nn.Module):

    def __init__(self, in_channels, representation_size):
        super(TwoMLPHead, self).__init__()

        self.fc6 = nn.Linear(in_channels, representation_size)
        self.fc7 = nn.Linear(representation_size, representation_size)

    def forward(self, x):
        # x为经过roiAlign后
        # x shape: [batch_size*RPN输出的正负样本总数(512), out_channels, 7, 7]
        # flatten: -> [batch_size*RPN输出的正负样本总数(512), out_channels * 7 * 7]
        x = x.flatten(start_dim=1)

        x = F.relu(self.fc6(x))
        x = F.relu(self.fc7(x))

        return x


class FastRCNNPredictor(nn.Module):

    def __init__(self, in_channels, num_classes):
        super(FastRCNNPredictor, self).__init__()
        self.cls_score = nn.Linear(in_channels, num_classes)
        self.bbox_pred = nn.Linear(in_channels, num_classes * 4)

    def forward(self, x):
        # x shape: [batch_size*RPN输出的正负样本总数(512), 1024] 1024: 对应TwoMLPHead最终的输出维度
        # x = x.flatten(start_dim=1)
        scores = self.cls_score(x)
        bbox_deltas = self.bbox_pred(x)

        return scores, bbox_deltas


class FasterRCNN(FasterRCNNBase):
    # backbone 自行定义的特征提取网络
    # num_classes 检测的类别个数（需要加入背景类）eg: VOC: num_classes=21
    def __init__(self, backbone, num_classes=None,
                 # transform parameter
                 min_size=800, max_size=1333,      # 预处理resize时限制的最小尺寸与最大尺寸
                 image_mean=None, image_std=None,  # 预处理normalize时使用的均值和方差(ImageNet)
                 # RPN parameters
                 # rpn_anchor_generator用于生成anchor的生成器。
                 # rpn_head由一个3*3conv、分类层和边界框回归层组成
                 rpn_anchor_generator=None, rpn_head=None,
                 # NMS前后相同主要是针对带有FPN的网络。由于FPN存在多个预测的特征图，每层在NMS前proposal数都保留2000个，总共加起来就上万了，然后再通过NMS保留2000个。
                 # 通过预测信息和anchor生成器，可以得到一系列的proposal。根据预测的score，在进行NMS之前先滤除一部分proposal，将剩余的proposal输入到NMS。
                 rpn_pre_nms_top_n_train=2000, rpn_pre_nms_top_n_test=1000,    # rpn中在nms处理前保留的proposal数(根据score)
                 rpn_post_nms_top_n_train=2000, rpn_post_nms_top_n_test=1000,  # rpn中在nms处理后保留的proposal数
                 rpn_nms_thresh=0.7,  # rpn中进行nms处理时使用的iou阈值
                 # rpn计算损失时，采集正负样本设置的阈值。anchor与GT的IOU大于0.7标记为正样本，anchor与任何一个GT的IOU均小于0.3标记为负样本
                 # fg: foreground bg: background
                 rpn_fg_iou_thresh=0.7, rpn_bg_iou_thresh=0.3,
                 # 在正样本和负样本中进行随机采样  计算RPN损失时。总共采样256个样本。rpn_positive_fraction为正样本占全部样本的比例
                 rpn_batch_size_per_image=256, rpn_positive_fraction=0.5,  # rpn计算损失时采样的样本数，以及正样本占总样本的比例
                 # Box parameters（ROI Head中的参数）
                 # box_roi_pool对应ROI Pooling
                 # box_head 对应MLPHead
                 # box_predictor对应两个fc层，一个预测类别概率、另一个用于预测边界回归框参数。
                 box_roi_pool=None, box_head=None, box_predictor=None,
                 # 移除小概率目标的阈值    fastRCNN中进行nms处理的阈值  对预测结果根据score排序取前100个目标
                 box_score_thresh=0.05, box_nms_thresh=0.5, box_detections_per_img=100,
                 # fastRCNN计算误差时，采集正负样本设置的阈值.proposal与GT的IOU大于0.5，定义为正样本。proposal与所有的GT的IOU均小于0.5，定义为负样本。
                 box_fg_iou_thresh=0.5, box_bg_iou_thresh=0.5,
                 # 共采样512个样本，正样本占全部样本的0.25.
                 box_batch_size_per_image=512, box_positive_fraction=0.25,  # fast rcnn计算误差时采样的样本数，以及正样本占所有样本的比例
                 bbox_reg_weights=None):

        # 预测特征层的channels resnet50+FPN out_channels=256
        out_channels = backbone.out_channels

        # 对数据进行标准化，缩放，打包成batch等处理部分（预处理）
        # 预处理的图像均值和方差
        if image_mean is None:
            image_mean = [0.485, 0.456, 0.406]
        if image_std is None:
            image_std = [0.229, 0.224, 0.225]

        transform = RcnnTransforms(min_size, max_size, image_mean, image_std)

        # 若anchor生成器为空，则自动生成针对resnet50_fpn的anchor生成器
        if rpn_anchor_generator is None:
            anchor_sizes = ((32,), (64,), (128,), (256,), (512,))
            aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes) # aspect_ratios将(0.5, 1.0, 2.0)重复五次，每一个元素对应一个特征层上的尺度
            rpn_anchor_generator = AnchorsGenerator(
                anchor_sizes, aspect_ratios
            )

        # 生成RPN通过滑动窗口预测网络部分
        if rpn_head is None:
            # rpn_anchor_generator.num_anchors_per_location()[0]: 对应每一层预测特征图上生成anchor的数量 由size和ratios决定
            rpn_head = RPNHead(
                out_channels, rpn_anchor_generator.num_anchors_per_location()[0]
            )

        # 默认rpn_pre_nms_top_n_train = 2000, rpn_pre_nms_top_n_test = 1000,
        # 默认rpn_post_nms_top_n_train = 2000, rpn_post_nms_top_n_test = 1000,
        rpn_pre_nms_top_n = dict(training=rpn_pre_nms_top_n_train, testing=rpn_pre_nms_top_n_test)
        rpn_post_nms_top_n = dict(training=rpn_post_nms_top_n_train, testing=rpn_post_nms_top_n_test)

        # 定义整个RPN框架
        # rpn_batch_size_per_image 为RPN计算损失时采用的正负样本的总个数
        # rpn_positive_fraction 在计算损失时 正样本占全部样本的比例
        # rpn_pre_nms_top_n 在进入NMS之前 对于每一个预测特征图所保存的目标个数
        rpn = RegionProposalNetwork(
            rpn_anchor_generator, rpn_head,
            rpn_fg_iou_thresh, rpn_bg_iou_thresh,
            rpn_batch_size_per_image, rpn_positive_fraction,
            rpn_pre_nms_top_n, rpn_post_nms_top_n, rpn_nms_thresh)

        # ROIAlign Pooling
        # 经过roiAlign后，每一个proposal的特征矩阵为7*7
        if box_roi_pool is None:
            box_roi_pool = MultiScaleRoIAlign(
                featmap_names=['0', '1', '2', '3'],  # 在哪些特征层进行roi pooling
                output_size=[7, 7], # 指定输出大小
                sampling_ratio=2)

        # FastRCNN中ROI-pooling后的展平处理两个全连接层部分（MLPHead）
        if box_head is None:
            resolution = box_roi_pool.output_size[0]  # 默认等于7
            representation_size = 1024
            box_head = TwoMLPHead(
                out_channels * resolution ** 2,
                representation_size
            )
        # 在box_head的输出上预测部分
        if box_predictor is None:
            representation_size = 1024
            box_predictor = FastRCNNPredictor(
                representation_size,
                num_classes)

        # 将roi pooling, box_head以及box_predictor结合在一起,组成RoIHead
        roi_heads = RoIHead(
            # box
            box_roi_pool, box_head, box_predictor, # 分别对应上面三个部分
            box_fg_iou_thresh, box_bg_iou_thresh,  # 0.5  0.5
            box_batch_size_per_image, box_positive_fraction,  # box_batch_size_per_image每张图像中选择的proposal的数量（512）  0.25
            bbox_reg_weights, # 默认(10 10 5 5 )
             # 对得到的最终结果进行后处理时使用  0.05  0.5  100
            box_score_thresh, box_nms_thresh, box_detections_per_img) 

        super(FasterRCNN, self).__init__(backbone, rpn, roi_heads, transform)