OneNet: Towards End-to-End One-Stage Object Detection 论文学习

Abstract

截至目前，端到端的单阶段目标检测仍落后于其它方法。本文研究发现，在标签分配过程中，候选样本与 ground-truth 之间缺少分类损失，这是单阶段检测器无法去除 NMS、实现端到端的主要障碍。现有单阶段目标检测器仅通过定位损失来赋予标签，比如 box IOU 或者点距离。不靠分类损失，在推理过程中，单凭定位损失会得到许多冗余的高置信度边框，这样NMS后处理就非常必要了。

为了设计一个端到端的单阶段目标检测器，作者提出了最小损失分配（Minimum Cost Assignment）。该损失是候选样本与 ground-truth 之间的分类损失和定位损失之和。对于每个 ground-truth 目标，只有损失最低的那个样本会被标注为正样本，其它的都是负样本。为了评价该方法的有效性，作者设计了一个极其简单的单阶段检测器，叫做 OneNet。结果表明，用最小损失分配训练出的 OneNet 模型，不会产生大量的重复边框，实现了端到端的检测。在COCO数据集上，OneNet 取得了 35.0 AP / 80 FPS，和 37.7 AP / 50 FPS，当输入分辨率是512像素。作者希望 OneNet 可以作为端到端单阶段模型的基线模型。代码位于 https://github.com/PeizeSun/OneNet。

1. Introduction

目标检测是计算机视觉领域的基础任务之一，使许多应用成为可能。它旨在图像中定位出物体，识别出其类别。当前目标检测器的一大挑战是标签分配问题。如何定义代表目标的正样本和代表背景的负样本一直是个开放问题。

几十年来，目标检测中的正样本就是与 ground-truth 边框的 IOU 大于某阈值的候选边框。在传统的计算机视觉中，人们在图像网格上滑动分类器。现在的检测器则会在图像网格上，预先定义好成百上千个 anchor boxes。我们称这些基于边框的标签分配方法为“边框分配”，如图2a所示。

尽管多年来边框候选方法统治着目标检测任务，但检测性能对目标大小、宽高比以及 anchor boxes 的数量很敏感。为了省去人为设计和复杂的边框候选的计算，有人提出了 anchor-free 检测器。这些方法直接将特征图上的网格点看作候选目标，然后预测网格点到目标边框边界的偏移量。Anchor-free 检测器中的标签分配问题就简化为了从网格点到目标边框中心点的点距离问题。我们将这种基于点的标签分配方法称为“点分配”，如图2b所示。

边框分配和点分配都要面对一个常见的困境，就是多对一的分配问题。对于单个 ground-truth 边框，它会有多于一个的正样本。在标签分配过程中，检测性能对于超参数是很敏感的。更糟糕的是，它会产生冗余、重复的预测结果，因此NMS后处理就很必要了。

图2：不同标签分配方法的比较。$H$和$W$表示特征图的高度和宽度，$K$是目标类别的个数。以前的单阶段目标检测方法只通过定位损失来赋予标签，如样本与 ground-truth 之间的 (a) IOU，(b) 点距离。在本文方法中，作者引入了 © 分类损失。作者发现，分类损失是端到端方法成功的关键因素。没有分类损失，仅凭定位损失会在推理过程中，会造成许多冗余的高置信度边框，使得NMS后处理非常必要。

最近在稀疏候选框和多阶段优化检测器中，一对一的标签分配取得了巨大成功，就是单个 ground-truth 边框只会被分配一个正样本，其它都是负样本。它们直接输出检测结果，而无需NMS。这些检测器的结果非常有前景。但是，密集候选框和单阶段优化检测器会更加简洁、快速。这就促使作者去设计一个直接的、一对一的标签分配策略，用于端到端的单阶段目标检测器。

一个直观的想法就是，我们可以将以前单阶段检测器中的多对一的标签分配替换为一对一的，比如只有一个正样本会被分配给一个 ground-truth 边框，特征图上其它的样本都是负样本。但是，这样做性能会很糟糕，如表1所示。在推理过程中，对于单个目标，检测器仍然输出多个高置信度的冗余边框，因而我们仍然需要 NMS。

作者发现对于单阶段目标检测器，要想省去NMS，真正地实现端到端，其主要障碍是在标签分配过程中，样本和 ground-truth 之间缺少分类损失。通过调研现有的单阶段目标检测器，作者发现它们只用了定位损失来分配标签，如边框 IOU 或点距离。但是，效果不错的目标检测器都是通过定位损失和分类损失来分配标签的。实验结果显示，推理过程中如果没有分类损失，就会导致冗余的高置信度边框，NMS因此就很必要。

为了设计一个端到端的单阶段目标检测器，作者提出了一个简单但有效的标签分配策略，叫做最低损失分配（Minimum cost assignment），其中损失是分类损失和定位损失的和。对于每个 ground-truth 边框，只有一个损失最低的样本会被标注为正样本；其它的都是负样本，如图2c所示。为了证明该策略的有效性，作者提出了 OneNet，它是单阶段的，且只需一个正样本。实验结果显示，通过最低成本分配策略训练出的 OneNet 避免了重复边框，实现了端到端的目标检测器。OneNet 展现了多个优势：

• 整个网络都是全卷积的，端到端的训练。过程中不需要 ROI 操作或注意力交互。
• 标签分配策略基于分类和定位的最低损失，而非人为设计的启发式规则或复杂的相互匹配。
• 无需任何后处理，如NMS或最大池化，使得目标检测更加高效率。

在COCO数据集上，图像大小是512像素时，OneNet 使用ResNet50 实现了 35.0 AP / 80 FPS，使用 ResNet101 实现了 37.7 AP / 50 FPS。作者认为更大的图像分辨率以及更复杂的主干网络能进一步提升检测性能，但是推理速度会有所降低，如表3和表5所示。该方法可以作为端到端单阶段目标检测的有效基线模型。因此，本文所报告的性能是在准确率和速度之间做了取舍。

通过定位损失来分配标签非常直观，在目前的目标检测方法中很流行。但是作者惊人地发现，该方法是端到端目标检测的障碍。作者解释到，主要原因是标签分配策略与网络优化目标的不一致。目标检测是分类与定位的多任务混合体。只通过定位损失来选择正样本，能最大化定位任务的贡献，但是无法保证最佳的分类表现。最佳的方案就是，一个目标只有一个预测。但是，如果我们只用定位损失来训练模型，分类分支就被强迫输出相似的结果；单个目标就会有多个预测。作者希望本文能够促使人们去重新思考目标检测中标签分配的问题，探索出下一代的目标检测器。

2. Related Work

目标检测。目标检测是计算机视觉领域最基础且具有挑战性的任务之一。受传统特征提取方法的限制，其表现几十年来都没什么进展，应用领域也有限。伴随着深度学习的迅速发展，目标检测实现了优异的性能。当前的目标检测器主要分为单阶段和双阶段。

单阶段目标检测器。单阶段目标检测器在不同的空间位置和尺度上，一次性地直接预测密集的 anchor boxes 或点的类别和位置，如 YOLO、SSD 和 RetinaNet。YOLO 将图像划分为 $S\times S$个网格，如果某物体的中心落入一个网格内，则此网格负责检测该物体。SSD 在卷积网络中的多尺度特征图层上铺放 anchor boxes，直接预测目标类别和 anchor box 偏移量。RetinaNet 利用 focal loss 来缓解正负样本的极端不均衡，基于 FPN。最近有人提出了 anchor-free 算法，使得这个流程更加简洁，用点来代替人为设计的 anchor boxes。上述方法都构建在密集点之上，直接对每个点做分类和回归。根据预先定义好的原则，比如 anchor 是否与 ground-truth 有着较高的 IOU，或者引用点是否落入目标框内，在训练时这些点就被赋给 ground-truth 目标框。CornerNet 通过热力图会生成上千个关键点，然后用 Associative Embedding 将它们分组，CornerNet 的表现非常优异，但是 false positives 过多。CenterNet 在单个尺度上，直接用中心点来回归目标物体。FCOS 将不同大小和尺度的物体分配到多尺度特征图上，有了强大的 FPN，FCOS 可以取得 SOTA 的性能。ATSS 揭示了基于 anchor 和 anchor-free 检测器的本质区别，就是如何定义正负训练样本，这导致了二者间的差异。尽管取得了显著成绩，在单阶段密集检测器推理时，我们仍然需要 NMS 后处理来去除多余的预测。

双阶段检测器。双阶段目标检测器是另一个优秀的方法，多年来引领着目标检测。该方法可以看作为单阶段目标检测器的扩展。它首先通过 RPN 网络，从密集的区域 anchors 中产生一组高质量的前景候选框，然后优化每个候选框的位置，预测其类别。该区域候选框算法在第一阶段中扮演着核心角色。Fast R-CNN 利用 Selective Search 来产生前景候选框，在 R-CNN 头中优化候选框。Faster R-CNN 提出了 RPN 网络，通过 CNN 网络产生一组高质量的候选框。Cascade RPN 系统性地解决了传统 RPN 的束缚，改进了区域候选框的质量和检测性能，启发式地定义 anchors，并将特征与 anchors 对齐。最近 CPNDet 提出了一个新的角点候选网络，替代 RPN。CPNDet 可以产生更佳的候选框，可以看作为双阶段 anchor-free 目标检测方法。在优化阶段，针对前景和背景预先定义的采样方法就很重要。Cascade R-CNN 迭代地利用多个 R-CNN 头，它们的标签分配阈值各不相同，从而得到高质量的检测框。Libra-RCNN 尝试解决采样层级、特征层级和目标层级的不均衡问题。Grid R-CNN 采用一个网格引导定位机制来实现准确的目标检测，而非传统的边框回归。这些双阶段方法都有很高的检测准确率。但是在候选区域阶段和目标识别阶段，它们都包含了过多的超参数和模块；因此，双阶段检测器在工业界应用起来不容易。

端到端的目标检测。优异的端到端的目标检测器都是基于稀疏候选框和多阶段优化步骤而来。DETR 直接输出检测结果，无需任何人工标签分配和后处理步骤，取得了梦幻般的效果。DETR 可以看作为第一个端到端的目标检测方法；DETR 使用一个稀疏的目标候选集合，与全局图像特征作交互。得益于全局注意力机制以及预测框和 ground-truth 目标之间的相互匹配，DETR 可以放弃 NMS 步骤，而取得杰出的表现。Deformable-DETR 将每个目标候选约束为一个由关键采样点组成的小集合，而非特征图上所有的点。Sparse R-CNN 始于一个候选目标框的固定稀疏集合，然后在目标识别头中迭代地进行分类和定位。

3. 标签分配

最低损失分配用于解决端到端目标检测过程中的标签分配问题。作者首先介绍了以前的单阶段检测器的标签分配方法，如 Box Assignment 和 Point Assignment。

3.1 匹配损失

以前的方法通过样本与 ground-truth 之间的边框 IOU 或点距离来分配样本。它们通过定位损失来分配标签。该定位损失的定义如下：

$$C_{loc}={\lambda }_{iou}\cdot C_{iou}+{\lambda }_{L_1}\cdot C_{L_1}$$

其中 $C_{L_1}$ 和 $C_{iou}$ 是样本和 ground-truth之间的 L1 损失和 IOU 损失。${\lambda }_{L_1}$ 和 ${\lambda }_{iou}$ 是系数。在 Box Assignment 中，${\lambda }_{L_1}=0$。在 Point Assignment 中，${\lambda }_{iou}=0$。

通过定位损失来分配标签在概念上很直观，检测性能也很优异。但是，目标检测是定位和分类的混合任务。只有定位损失会造成分类性能不是最佳。它会造成过多冗余的高置信度的边框，不得不需要 NMS 后处理。在第五部分会有更多详细的讨论。

作者在匹配损失中引入了分类损失，从而实现端到端的单阶段目标检测。该损失是样本与 ground-truth 之间分类损失和定位损失的和，定义如下：

$$C={\lambda }_{cls}\cdot C_{cls}+C_{loc}$$

其中$C_{cls}$是预测结果与 ground-truth 类别标签之间的分类损失。$C_{loc}$ 在等式一中有定义。${\lambda }_{cls}$ 是系数。在 Box Assignment 和 Point Assignment 中，${\lambda }_{cls}=0$。

3.2 Minimum Cost Assignment

最低损失分配非常直接：对于每个 ground-truth，只有最低损失的样本才会被选做正样本；其它的都是负样本。这里面没有涉及任何的人为启发式规则或复杂的相互匹配。算法一是最低损失分配的例子：分类损失用交叉熵损失，定位损失则是 L1 损失。

在密集检测器中，分类损失是 Focal Loss。定位损失包含 L1 损失和 generalized IOU（GIOU） 损失。最后，损失定义如下：

$$C={\lambda }_{cls}\cdot C_{cls}+{\lambda }_{L_1}\cdot C_{L_1}+{\lambda }_{giou}\cdot C_{giou}$$

其中$C_{cls}$是 focal loss，$C_{L_1}$和$C_{giou}$是预测边框的归一化中心坐标和高度及宽度与 ground truth 边框的 L1 损失和 GIOU 损失。${\lambda }_{cls}$、${\lambda }_{L_1}$和${\lambda }_{giou}$是系数。

作者注意到[41]论文中，对每个 ground-truth 只选择一个正样本，而附近的样本则用高斯核来标注。本文方法则更加直接：除了选中的正样本，其它的都是负样本。

4. OneNet

OneNet 是一个全卷积单阶段检测器，无需后处理步骤如NMS。整个流程如图3所示。

主干网络。主干网络首先是一个自下而上的结构，然后自上而下结构。自下而上的结构是一个 ResNet，产生多尺度特征图。自上而下的结构搭配FPN侧连接，用于产生最终的特征图，进行目标识别。输出特征的形状是 $H/4+W/4\times C$，其中$H$和$W$是图像的高度和宽度。

头部。通过两个平行的卷积层，在$H/4\times W/4$的特征图中的每个网格点上，头部进行分类和定位。分类层预测物体出现在每个网格点的概率，共有$K$个目标类别。定位层预测每个网格点到 ground-truth 边框四条边的偏移量。

训练。标签分配策略是最低损失分配。训练损失与匹配损失相似，由 focal loss、L1 损失和GIOU损失组成。

推理。最终的输出直接就是 top-k 个得分边框。无需任何的后处理，如NMS或最大池化操作。

4.1 Multi-head Training

作者提出了一个可选的 Multi-head 训练策略。它主要由级联预测头部和共享权值机制组成。

级联头部。对于第一个阶段，输入特征表示为 $F_0$，它的维度通过 broadcast 增加一倍，$H/4\times W/4\times C\to H/4\times W/4\times 2C$。然后它输入进一个卷积层，输出特征 $F_1$的形状是$H/4\times W/4\times C$。基于 $F_1$，分类预测和回归预测就得到了。对于后面的阶段$j$，原始特征$F_0$的形状是 $H/4\times W/4\times C$，而前面阶段的特征 $F_{j-1}$ 的形状也是$H/4\times W/4\times C$，将二者沿着通道维度 concat 起来，得到了形状为$H/4\times W/4\times 2C$的复合特征。然后就产生了形状是$H/4\times W/4\times C$的特征图$F_j$，进行分类和回归。

权值共享。在每个头部中，分类卷积和回归卷积共享权值。

仅使用 Multi-head 训练不会带来任何检测准确率的提升，反而推理速度会下降。针对 Multi-head 训练，作者提出了两点改进，使最终的性能获益：大的学习率和 Single-head 推理。

大学习率。较大的学习率能够带来准确率提升。但是直接增加单个预测头部的学习率会使得训练不稳定，使检测准确率退化。当使用级联头部和权值共享，训练学习率就可以增大，提升准确率。

Single-head 推理。推理时，只有第一个阶段会输出最终的结果，其它阶段不用。这造成的准确率下降微乎其微。所以，不会增加任何的额外计算开支，而准确率表现得到提升。Multi-head 训练是 OneNet 可选的训练策略。实验结果显示，在 COCO 数据集上，三个阶段足以提升 1 AP， 如表2所示。

5. 实验