OneNet:无需NMS的One-stage端到端目标检测方法.

背景

目前的一些目标检测方法还不够end-to-end,通常需要对生成的冗余检测框进行NMS,其原因是在模型训练时,生成这些检测框时进行了标签匹配(label assignment)。标签匹配是指建立生成检测框与标注Ground Truth框之间的匹配关系,并构造匹配损失作为优化目标的一部分。

主流的目标检测方法使用的匹配方式是基于IoU的,即当候选框与Ground Truth框的IoU超过一定阈值则认为其是正样本。这是一种框匹配(box assignment)的方式。框匹配时通常人为预设大量anchor,这些是anchor-based的方法。

也有一些anchor-free的方法,不需要人为预设anchor框,而是把网格点看作样本,根据网格点与目标点的距离和网格点是否在目标内部来判断其是否是正样本。这是一种点匹配(point assignment)的方式。

上述两种匹配方式共同的问题是“多对一”,即每一个Ground Truth通常有多个正样本与之匹配。这样容易产生冗余的检测结果,后续处理时必须使用NMS。作者提出了一种简单的匹配方法:minimum cost assignment。在进行标签匹配时,引入分类损失,对每个Ground Truth只有一个具有最小分类损失和定位损失的样本被分配为正样本。

上述三种分配损失如下图所示:

标签分配:minimum cost assignment

之前的方法是通过IoU或坐标点的距离来进行匹配的,作者将其总结为位置损失。其定义如下:

\[\mathcal{C}_{loc} = \lambda_{iou} \cdot \mathcal{C}_{iou} + \lambda_{L1} \cdot \mathcal{C}_{L1}\]

其中\(\mathcal{C}_{iou}\)是IoU损失,\(\mathcal{C}_{L1}\)是L1损失。对于框分配,$\lambda_{L1} = 0$;对于点分配,$\lambda_{iou} = 0$。

目标检测是一个多任务,既有回归定位又有分类。只使用位置损失会导致高置信度的冗余框产生,导致后续处理需要NMS。作者将分类损失也引入匹配损失,定义如下:

\[\mathcal{C} = \lambda_{cls} \cdot \mathcal{C}_{cls} + \mathcal{C}_{loc} = \lambda_{cls} \cdot \mathcal{C}_{cls} + \lambda_{iou} \cdot \mathcal{C}_{iou} + \lambda_{L1} \cdot \mathcal{C}_{L1}\]

下面Pytorch形式的伪代码给出了一个minimum cost assignment的例子。其中类别损失采用交叉熵损失,定位损失采用L1损失。

# For simplicity,
# cross entropy loss as classification cost
# L1 loss as location cost

# Input:
# class_pred, box_pred: network output(HxWxK, HxWx4)
# gt_label, gt_box: ground-truth (N, Nx4)

# Output:
# src_ind: index of positive sample in HW sequence(N)
# tgt_ind: index of corresponding ground-truth (N)

# flattened class: HWxK
output_class = class_pred.view(-1, K)

# flattened box: HWx4
output_box = box_pred.view(-1, 4)

# classification cost: HWxN
cost_class = -torch.log(output_class[:, gt_label])

# location cost: HWxN
cost_loc = torch.cdist(out_box, gt_bbox, p=1)

# cost matrix: HWxN
cost_mat = cost_class + cost_loc

# index of positive sample: N
_, src_ind = torch.min(cost_mat, dim=0)

# index of ground-truth: N
tgt_ind = torch.arange(N)

模型

模型的主要流程如上图所示。其主要结构如下:

作者还提出了一种multi-head训练策略。训练时串联多个预测head,其中的分类和回归网络共享参数;推理时只使用第一个head。整体结构如下:

输入特征图用$F_0$表示,将$F_0$的通道维度广播为两倍,得到维度为$\frac{H}{4} \times \frac{W}{4} \times 2C$。 然后通过卷积得到$F_1$,$F_1$的维度为$\frac{H}{4} \times \frac{W}{4} \times C$。 然后基于$F_1$进行分类和回归。 对于后面的阶段$j$,原始的特征$F_0$和前一个阶段的特征$F_{j-1}$拼接起来,得到的维度为$\frac{H}{4} \times \frac{W}{4} \times 2C$,然后再卷积得到$F_j$,维度为$\frac{H}{4} \times \frac{W}{4} \times C$,然后基于$F_j$进行分类和回归。

单纯引入这个多阶段训练对精度提升不大,还会影响训练速度。之所以这么做是为了使用更大的学习率和单阶段推理。

实验

作者对CenterNetOneNet的正样本进行了可视化,如下图所示。CenterNet中正样本都位于目标框的中心位置,这对框的回归是有利的,但对正负样本的区分并不是最好的。OneNet中正样本定位到人体上特征区分最明显的部分。