BorderDet 原理与代码解析

paper：BorderDet: Border Feature for Dense Object Detection

code：https://github.com/Megvii-BaseDetection/BorderDet

背景

大多数滑动窗口形式的目标检测模型都是用输出特征图上的单点特征去预测目标的分类和回归，但是受有限感受野的影响单点特征可能没有包含足够的信息来表示一个完整的目标，同时它可能也缺乏目标边界的信息来精确的回归边界框。许多研究都集中在物体的特征表示上，如下图所示

这些方法比单点特征提取了更多的特征，但应用于密集目标检测时存在两点限制：

1. 从整个box内提取的特征可能包含不必要的背景信息。

2. 这些方法都是隐式、间接地提取边界特征的，由于特征是在整个box内自适应的提取的，这些方法并没有专门提取边界特征。

本文的创新点

作者提出了一个新的特征提取方法BorderAlign，它直接利用从每个边界上提取的特征作为原始单点特征的补充，增强提取的特征信息。和图1中其它从整个box内密集地提取特征的方法不同，BorderAlign聚焦于物体的边界，并且自适应地提取边界上最关键部分的特征，比如目标的端点，如图1(a)中的四个橘色原点所示。

作者设计了一个新的检测模型BorderDet，其中利用边界对齐模块Border Alignment Modules（BAM）来细化分类得分和边界回归。和类似的特征增强方法相比，BorderDet用更少的计算量获得了更高的精度，并且可以方便地集成到其它检测模型中。

方法介绍

作者首先做了个对比实验比较bounding box各种特征表示的精度差异，首先采用一个简单的检测器FCOS作为baseline来生成coarse边界框预测，然后如下图2中(b)-(e)的方法重新提取特征，然后逐步用不同的特征作为单点特征的补充来refine粗边界框预测结果。

实验结果如下表所示

根据对比实验的结果，作者得到了以下结论：

1. 区域特征比单点特征的表示能力更强。用区域特征作为补充特征AP提升了1.3。

2. 当使用区域特征来增强单点特征时，边界特征在区域特征中起主要作用。当忽略区域内部，只引入边界特征时，AP只下降了0.3。

3. 有效地提取边界特征比密集地提取边界特征能进一步提升模型精度。提取每条边界中间点的特征比密集地提取边界上的特征，AP提升了0.3，用更少的采样点得到了和提取区域特征一样的精度。

Border Align

基于上述观察，边界特征对于获得更好的检测性能非常重要，由于边界上前景很少同时存在大量的背景，在边界上密集的提取特征是低效的。因此作者提出了一种新的特征提取方法BorderAlign来高效的利用边界特征。

BorderAlign的结构如上图3所示，其输入为通道数为5C的border-sensitive特征图，其中4C对应左、上、右、下四条边界的特征，另外一C对应原始的单点特征。然后每条边被均匀地分成N个点，文中N=10，对N个点对应的特征进行max-pooling，由于N个点的坐标可能为小数，因此每个点的特征通过双线性插值得到。这样最终每条边会输出一个采样点的特征，原始一个单点anchor point会采样1个原始点的特征加上4条边上各一个点的特征共5个点的特征。

需要注意的是在对一条边上的N个点进行最大池化时是在对应的C通道内进行的，假设输入特征图的5C通道顺序为单点、左边、上边、右边、下边，则输出特征图 \(\mathcal{F}\) 可以表示为下式

其中 \(\mathcal{F}_{c}(i,j)\) 表示输出特征图 \(\mathcal{F}\) 中第 \(c\) 个通道 \((i,j)\) 处的值，\((x_{0},y_{0},x_{1},y_{1})\) 是点 \((i,j)\) 处预测的bounding box坐标，\(w,h\) 是预测的bouding box的宽高，\(I_{c}\) 的值通过双线性插值得到。

对特征图每个C通道上的最大值进行可视化结果如下，可以看出每条边上的前景区域被引导响应。

BorderDet

BorderDet的完整结构如图3所示，采用FCOS作为baseline，因为BorderAlign需要边界位置作为输入，金字塔特征图即FPN的输出作为输入，首先采用FCOS的预测作为coarse分类得分预测和边界位置预测，然后将FPN对应层的输出特征图和coarse边界预测输入到BAM中得到包含显式边界信息的输出特征图，然后接1x1卷积得到border分类得分和边框预测，最终统一两个阶段的预测作为最终预测结果。

Border Alignment Module

BAM的结构如图3所示，输入特征图通道数为C，通过1x1卷积和instance normalization得到border-sensitive特征图，通道数为5C，文中分类分支C=256，回归分支C=128。BorderAlign从border-sensitive特征图中提取了边界特征得到输出特征图，通道数依然为5C，最终通过1x1卷积将通道数降为C。

Model Training and Inference

Target Assignment

第一阶段FCOS的分类和回归预测的target和原始FCOS中一样，第二阶段border分类和回归target的分配如下：第一阶段的边框预测结果 \((x_{0},y_{0},x_{1},y_{1})\) 作为第二阶段的输入，通过IoU阈值0.6将其分配给对应的gt box \((x_{0}^{t},y_{0}^{t},x_{1}^{t},y_{1}^{t})\)，对应的回归target \((\delta x_{0},\delta y_{0}, \delta x_{1}, \delta y_{2})\) 按下式计算

其中 \(w,h\) 是第一阶段coarse预测边界框的宽高，\(\sigma\) 是超参文中设置为0.5。

Loss Function

多任务损失函数如下所示

其中 \(\mathcal{L}^{C}_{cls}\) 和 \(\mathcal{L}^{C}_{reg}\) 分别是第一阶段coarse预测的分类损失和回归损失，文中分别采用Focal Loss和IoU Loss。\(\mathcal{L}^{B}_{cls}\) 是第二阶段border分类得分和为其分配的gt \(\mathcal{C}^{*}\) 之间的Focal Loss，除以正样本数量 \(\mathcal{N}_{pos}\)。第二阶段border的回归损失 \(\mathcal{L}^{B}_{reg}\) 采用 \(\mathcal{L}_{1}\) loss。\(\mathcal{P}^{B}\) 表示第二阶段预测的border分类得分，\(\triangle\) 是第二阶段预测的边界偏移。

Inference

注意，在训练阶段第一阶段的分类、回归和第二阶段的分类、回归是分开的四个分支，分别单独优化的。但在推理阶段，第一阶段的分类得分和第二阶段的分类得分相乘得到最终的分类得分，对应图3中的操作\(\pi\)。第一阶段的边框预测和第二阶段的边框预测相加得到最终的边框预测结果，对应图3中的操作 \(\delta\)。

实验结果

消融实验

本文的最大创新点就在于BAM，作者将BAM分别添加到分类和回归分支，以及同时添加进去，对比了最终精度，结果如下。可以看出，两个分支分别加入BAM精度都得到了提升，同时加入BAM精度进一步得到了提升。

Border Align

Pooling Size

Border Align中对每条边上N个点进行max pooling得到最终采样点特征，N太小时结果不稳定，N太大时增加计算量，从下表可以看出当N=10时，精度最高。

Border-Sensitive Feature Maps

文中BAM的输入是通道数为5C的border-sensitive特征图，即每C个通道对应不同的边或原始anchor point，作者这里对比了采用通道数为C的border-agnostic特征图作为输入，即anchor point和4条边的特征都从这C个通道中进行提取，结果如下。可以看出border-sensitive的精度更高。

Border Feature Aggregation Strategy

原始BorderAlign中采用的是channel-wise的max pooling策略，即对某条边的C通道进行max pooling时，最大池化是在每一个通道内单独进行的。这里作者对比了border-wise的策略，即沿所有通道用avg pooling或max pooling先进行池化，这样每条边就只对应单通道的特征图，然后再沿边的N各点进行max pooling，对比结果如下，可以看到channel-wise的精度更高。

Comparision with SOTA

和当时的其它SOTA方法的对比如下，可以看出BorderDet获得了最高的精度。

代码解析

代码实现文件在https://github.com/Megvii-BaseDetection/BorderDet/blob/master/playground/detection/coco/borderdet/borderdet.res50.fpn.coco.800size.1x/borderdet.py

输入batch_size=2，预处理后input_shape=(2, 3, 1088, 800)，然后提取FPN汇总p3-p7层的输出作为head的输入，代码如下

features = self.backbone(images.tensor)  # (2,3,1088,800)
features = [features[f] for f in self.in_features]  # ['p3', 'p4', 'p5', 'p6', 'p7']
# [(2,256,136,100),(2,256,68,50),(2,256,34,25),(2,256,17,13),(2,256,9,7)]

本文的创新点都在head部分，self.head(features, shifts)调用head，进入BorderHead类的forward函数中，实现代码如下

def forward(self, features, shifts):
    """
    Arguments:
        features (list[Tensor]): FPN feature map tensors in high to low resolution.
            Each tensor in the list correspond to different feature levels.

    Returns:
        logits (list[Tensor]): #lvl tensors, each has shape (N, K, Hi, Wi).
            The tensor predicts the classification probability
            at each spatial position for each of the K object classes.
        bbox_reg (list[Tensor]): #lvl tensors, each has shape (N, 4, Hi, Wi).
            The tensor predicts 4-vector (dl,dt,dr,db) box
            regression values for every shift. These values are the
            relative offset between the shift and the ground truth box.
        centerness (list[Tensor]): #lvl tensors, each has shape (N, 1, Hi, Wi).
            The tensor predicts the centerness at each spatial position.
        border_logits (list[Tensor]): #lvl tensors, each has shape (N, K, Hi, Wi).
            The tensor predicts the border classification probability
            at each spatial position for each of the K object classes.
        border_bbox_reg (list[Tensor]): #lvl tensors, each has shape (N, 4, Hi, Wi).
            The tensor predicts 4-vector (dl,dt,dr,db) box
            regression values for every border shift. These values are the
            relative offset between the shift and the ground truth box.
        pre_bbox (list[Tensor]): #lvl tensors, each has shape (N, Hi * Wi, 4).
            The tensor predicts 4-vector (l,t,r,b) box regression values.
            These values are predicted boxes by the dense object detector.
    """
    logits = []
    bbox_reg = []
    centerness = []
    border_logits = []
    border_bbox_reg = []
    pre_bbox = []

    # [[(13600,2),(3400,2),(850,2),(221,2),(63,2)],
    #  [(13600,2),(3400,2),(850,2),(221,2),(63,2)]]
    # list(zip(*shifts)
    # [((13600,2),(13600,2)),
    #  ((3400,2),(3400,2)),
    #  ((850,2),(850,2)),
    #  ((221,2),(221,2)),
    #  ((63,2),(63,2))]

    shifts = [
        torch.cat([shi.unsqueeze(0) for shi in shift], dim=0)
        for shift in list(zip(*shifts))
    ]  # [(2,13600,2),(2,3400,2),(2,850,2),(2,221,2),(2,63,2)]

    for level, (feature, shifts_i) in enumerate(zip(features, shifts)):
        # 0, (2,256,136,100), (2,13600,2)
        cls_subnet = self.cls_subnet(feature)  # (2,256,136,100)
        bbox_subnet = self.bbox_subnet(feature)  # (2,256,136,100)

        logits.append(self.cls_score(cls_subnet))  # (2,20,136,100)
        if self.centerness_on_reg:  # True
            centerness.append(self.centerness(bbox_subnet))  # (2,1,136,100)
        else:
            centerness.append(self.centerness(cls_subnet))
        bbox_pred = self.scales[level](self.bbox_pred(bbox_subnet))  # (2,4,136,100)
        if self.norm_reg_targets:  # True
            bbox_pred = F.relu(bbox_pred) * self.fpn_strides[level]
        else:
            bbox_pred = torch.exp(bbox_pred) * self.fpn_strides[level]
        bbox_reg.append(bbox_pred)

        # border
        N, C, H, W = feature.shape
        pre_off = bbox_pred.clone().detach()
        with torch.no_grad():
            pre_off = pre_off.permute(0, 2, 3, 1).reshape(N, -1, 4)  # (2,4,136,100)->(2,136,100,4)->(2,13600,4)
            pre_boxes = self.compute_bbox(shifts_i, pre_off)  # (2,13600,4), 在原图上预测的bbox
            align_boxes, wh = self.compute_border(pre_boxes, level, H,
                                                  W)  # (2,13600,4),(2,13600,2), 将原图上预测的bbox映射到某一level对应的feat_map上
            pre_bbox.append(pre_boxes)
        border_cls_conv = self.border_cls_subnet(cls_subnet, align_boxes, wh)  # (2,256,136,100)
        border_cls_logits = self.border_cls_score(border_cls_conv)  # (2,20,136,100)
        border_logits.append(border_cls_logits)

        border_reg_conv = self.border_bbox_subnet(bbox_subnet, align_boxes, wh)  # (2,256,136,100)
        border_bbox_pred = self.border_bbox_pred(border_reg_conv)  # (2,4,136,100)
        border_bbox_reg.append(border_bbox_pred)

    if self.training:
        pre_bbox = torch.cat(pre_bbox, dim=1)  # (2,18134,4), 18134=13600+3400+850+221+63
    return (logits, bbox_reg, centerness, border_logits, border_bbox_reg, pre_bbox)

其中self.cls_subnet和self.bbox_subnet分别是图3最左侧分类和回归分支的连续4个3x3-256卷积层。

self.cls_score和self.bbox_pred分别得到FCOS原始的分类得分和回归预测，对应图3中的Coarse Cls Score和Coarse Box Reg。self.centerness是原始FCOS独有的分支。

上面是原始FCOS的分类、回归、centerness三个分支的最终输出，即BorderDet中第一阶段的输出。下面是BorderDet新增的部分。

self.compute_bbox根据anchor point和回归分支预测结果即到四边的距离，得到原图上预测的所有box。self.compute_border是将原图上预测的box除以p3-p7层的stride将box映射到对应level的特征图上。

接下来self.border_cls_subnet和self.border_cls_score调用的都是BorderBranch类，分别是分类分支和回归分支中的Border Alignment Module。BorderBranch类的完整代码如下

class BorderBranch(nn.Module):
    def __init__(self, in_channels, border_channels):
        """
        :param in_channels:
        """
        super(BorderBranch, self).__init__()
        self.cur_point_conv = nn.Sequential(
            nn.Conv2d(
                in_channels,
                border_channels,
                kernel_size=1),
            nn.InstanceNorm2d(border_channels),
            nn.ReLU())

        self.ltrb_conv = nn.Sequential(
            nn.Conv2d(
                in_channels,
                border_channels * 4,
                kernel_size=1),
            nn.InstanceNorm2d(border_channels * 4),
            nn.ReLU())

        self.border_align = BorderAlign(pool_size=10)

        self.border_conv = nn.Sequential(
            nn.Conv2d(
                5 * border_channels,
                in_channels,
                kernel_size=1),
            nn.ReLU())

    def forward(self, feature, boxes, wh):  # (2,256,136,100),(2,13600,4)
        N, C, H, W = feature.shape

        fm_short = self.cur_point_conv(feature)  # (2,256,136,100)
        feature = self.ltrb_conv(feature)  # (2,1024,136,100)
        ltrb_conv = self.border_align(feature, boxes)  # (2,256,13600,4)
        ltrb_conv = ltrb_conv.permute(0, 3, 1, 2).reshape(N, -1, H, W)  # (2,256,13600,4)->(2,4,256,13600)->(2,1024,136,100)
        align_conv = torch.cat([ltrb_conv, fm_short], dim=1)  # (2,1280,136,100)
        align_conv = self.border_conv(align_conv)  # (2,256,136,100)
        return align_conv

forward函数中，首先self.cur_point_conv是1x1-256的卷积，self.ltrb_conv是1x1-1024的卷积得到通道数为4C的四条边界的border-sensitive特征图。self.border_align就是BorderAlign的具体过程，这里具体的实现在cvpods库中且是cuda的实现，这里就不放代码了，大致讲一下实现过程：输入为border-sensitive特征图和原始FCOS得到的Coarse Box Reg预测，将预测的box映射到对应level的border-sensitive特征图上，其通道数为4C，每C个通道对应一个边界，对C个通道中每个通道上box对应边界上均分的10个点通过双线性插值得到具体值，然后对10个值进行max pooling。self.border_align的输出shape=(2, 256, 13600, 4)，这里是大小为136x100的P3特征图上的四条边界上的4个采样点的特征，每个点对应256个通道的特征，回归分支中每个点对应128个通道的特征。然后将原始anchor point的特征和四个边界点的特征concat起来得到通道数为5C=5x256=1280的输出，最后再通过self.border_conv即一个1x1-256卷积将5C映射回C。

最后head中self.border_cls_score和self.border_bbox_pred得到border的分类和回归的最终输出结果。

在训练阶段，第一阶段FCOS原始的分类和回归以及第二阶段border的分类和回归四个输出是分别计算损失的。但在推理阶段，要将两个阶段的分类和回归分别融合得到最终的分类和回归输出。其中分类的融合在Line 526通过相乘进行融合

predicted_prob = predicted_prob * bd_box_cls_i

回归的融合在Line 546通过相加进行融合

predicted_boxes = bd_based_box_i + (bd_box_reg_i * border_bbox_std * det_wh)