从论文到源码实现—RPN

727 阅读2分钟

一起养成写作习惯!这是我参与「掘金日新计划 · 4 月更文挑战」的第20天,点击查看活动详情

我们将实现并理解两阶段物体检测的几个核心组件。说到两阶段物体检测,我们会想到 RCNN 系列模型 RCNN、Fast RCNN、Faster R-CNN 和 Mask RCNN。

所有的两级物体检测器都有几个主要的组成部分

  • backbone 网络: 用于特征提取的 CNN 网络结构
  • RPN: 生成可能候选区域,有时候也叫建议区域,也就是根据先验在图像或者图像经过 backbone 得到特征图上可能存在目标的区域
  • RoIPooling 和 RoIAlign:RPN 在特征图映射区域,大小不一致,为了输入到接下来用于分类和边界框回归,RoIPooling 和 RoIAlign 主要用于将这些大小不一致特征图保持,类似一个池化层,RoIPooling 是在 Fast RCNN 提出沿用到 Faster RCNN,不过 RoIPooling 因为存储两次取整而存在精度损失,RoIAlign 设计通过双线性插值避免精度问题。

RPN

RPN(Region Proposal Network) 候选提议网络,将 backbone 的最后卷积层输出作为 RPN 的输入。生成候选区域生成时用 3x3 来生成一个中间输出,这个中间输出和 backbone 输出形状保持一致,

rpn_001.PNG

代码实现

class RegionProposalNetwork(nn.Module):
    """
    input_size 输入
    layer_size 层的尺寸
    conv_size 卷积尺寸
    num_anchor anchor 数量
    """
    def __init__(self, input_size, layer_size, conv_size, num_anchor):
        super().__init__()

        self.input_size = input_size
        self.layer_size = layer_size
        self.num_anchor = num_anchor
        self.conv_size = conv_size

        self.intermediate = nn.Conv2d(
            self.input_size, self.layer_size, self.conv_size, stride=1, padding=1
        )
        #类别检测头
        self.classification_head = nn.Conv2d(self.layer_size, self.num_anchor, 1)
        #边界框回归头
        self.reggresion_head = nn.Conv2d(self.layer_size, 4 * self.num_anchor, 1)

        for layer in self.children():
            torch.nn.init.normal_(layer.weight, std=0.01)
            torch.nn.init.constant_(layer.bias, 0)

    def forward(self, feature_map):

        t = torch.nn.functional.relu(self.intermediate(feature_map))
        classification_op = self.classification_head(t)
        regression_op = self.reggresion_head(t)

        classification_op = classification_op.permute(0,2,3,1).flatten()
        regression_op = regression_op.permute(0,2,3,1).reshape(-1,4)

        return classification_op, regression_op

在网络中,定义一个中间层 intermediate 是一个卷积层,这是 3x3 卷积核,stride 为 1 和 padding 为 1 ,这样以来中间层 intermediate 输出 shape 和输入 shape 保持一致。self.classsification_head 是一个通过 1x1 卷积核将提取特征图转换为 2k 其中 k 为 anchor 数量,这里 2 表示对于 anchor 给出两 anchor 区域是前景还是背景的概率。另一个经过 self.reggresion_head 这个用于将提取特征图 anchor 转换 4k ,也就是 dx,dyd_x,d_y 是 anchor 中心点的偏移量,而 dw,dhd_w,d_h 是相对于 anchor 尺寸的缩放。然后分别对 classification_op 进行变换后展平。