作者丨happy
编辑丨极市平台
本文将单阶段检测器在COCO数据集上的性能刷新到了51.1AP。针对现有单阶段检测器分类与定位存在不对齐问题，本文提出了一种新颖T-Head对其进行平衡；同时还提出了TAL在训练过程中将两个任务的最优anchor进行显式靠拢(甚至统一)。受益于所提到的T-Head与TAL，所提方法TOOD刷新了COCO数据集上的单阶段检测器性能，达到了５１.１AP，超过了GFLv2、OTA、IQDet等方案。
pape: https://arxiv.org/abs/2108.07755
code: https://github.com/fcjian/TOOD
单阶段目标检测往往通过优化两个子任务进行实现：目标分类与目标定位，两者采用两个并行分支，这种处理方式可能会导致一定程度上的空间不对齐问题。
我们提出一种任务对齐单阶段目标检测(Task-aligned One-stge Object Detection ,TOOD)，它通过学习的方式对两个任务进行显式对齐。首先，我们设计了一种新颖的T-Head(Task-aligned Head)，它可以为任务交互以及任务相关特征提供更好的平衡，还可以通过任务对齐预测器为学习对齐提供更好的灵活性。其次，我们提出了TAL(Task Alignment Learning)在训练过程中通过采样分配机制与任务对齐损失将两个任务最终的anchor进行显式靠拢(甚至统一)。
我们在MS-COCO数据集上进行了充分了实验，凭借单模型单尺度测试，TOOD取得了51.1AP指标，大幅超过了已有单阶段检测(比如ATSS的47.7AP、GFL的48.2AP、PAA的49.0AP)，同时具有更少的参数量、FLOPs。此外，定量分析结果同样表明：TOOD可以更好的对目标分类与目标定位两个任务进行对齐。
本文主要贡献包含以下三点：设计了一个新的T-Head以增强分类与定位之间的相互作用，同时保持各自的特性，此外在预测过程中对两个任务进行对齐；提出了TAL对两个任务进行显式对齐，同时为所提预测器提供学习信号；所提TOOD在COCO数据集上取得了51.1AP指标，超越了包含ATSS、GFL、PAA、OTA等在内的单阶段检测。
类似于近期提出的单阶段检测器，所提TOOD采用了类似的架构：backbone-FPN-head。考虑到效率与简单性，类似ATSS， TOOD在每个位置放置一个anchor，即anchor-free。正如所讨论的，由于分类与定位任务的发散性，现有单阶段检测器存在任务不对齐(task misalignment)约束问题。本文提出通过显式方式采用T-head+TAL对两个任务进行对齐，见上图。T-head与TAL通过协同工作方式改善两个任务的对齐问题，具体来说，T-head首先在FPN特征基础上进行分类与定位预测；然后TAL基于所提任务对齐测度计算任务对齐信息；最后T-head根据从TAL传回的信息自动调整分类概率与定位预测。
我们的目标是设计一种高效Head结构以改善现有单阶段检测器中的Head(见上图a)，针对此，我们主要从以下两个角度进行考量：提升两个任务之间的相互作用；增强检测器学习对齐的能力。
本文所提T-Head见上图b，它具有非常简单的结构：特征提取器+TAP。为增强分类与定位之间的相互作用，我们通过特征提取器学习任务交互（task-interactive）特征，见Figure3b中蓝色框部分。这种设计不仅有助于任务交互，同时可以为两个任务提供多级多尺度特征。假设表示FPN特征，特征提取器采用N个连续卷积计算任务交互特征：
因此，通过特征提取器我们可以得到丰富的多尺度特征并用于送入到后续两个TAP模块中进行分类与定位对齐。
Task-aligned Predictor(TAP) 我们在前述所计算的任务交互特征基础上进行目标分类与定位，此时两个任务可以很好的感知到相互之间的状态信息。然而，由于单分支设计因素，任务相关特征不可避免地会引入一定程度地特征混淆。只管来讲，目标分类与目标定位具有完全不同地目标，因此聚焦于不同类型地特征(比如，不同层级、不同感受野)。因此，我们提出采用注意力机制进行任务分解，见上面Figure3c。此时所得特征我们称之为任务相关(task-specific)特征，描述如下：
注：这里的w采用类似SE注意力方式得到。最后，在上述特征基础上进行分类与定位：
Prediction alignment 在预测阶段，我们进一步对两个任务通过调整空间分布进行自适应对齐。不同于已有工作采用centerness分之或者IoU分支仅调整分类预测，我们对两个预测同时进行对齐，见Figure3c。从图示可以看到：我们采用空域概率图调整分类预测：
与此同时，为对定位预测进行对齐，我们进一步学习了空域移位图以调整每个位置预测的目标位置信息。具体调整方式如下：
通过上述对齐，我们可以进行更精确的目标定位。因此，所提方法不仅对两个任务进行了对齐，同时提升了定位精度。上面提到的两个对齐图学习方式如下：
注：M与O的学习是在TAL中实现，将在下面的内容进行介绍。此外需要注意的是：T-Head是一个独立模块，可以在无TAL的条件下工作。也即是说，T-Head可以通过“即插即用”方式快速应用到不同的单阶段检测器中并提升检测性能。
我们进一步引入TAL引导T-Head进行任务对齐预测。相比已有工作，TAL有以下两个不同：它是从任务对齐角度出发设计得到；它同时考虑了anchor分配与加权同步。
为与NMS搭配，训练样例的anchor分配需要满足以下规则：正常对齐的anchor应当可以预测高分类得分，同时具有精确定位；不对齐的anchor应当具有低分类得分，并在NMS阶段被抑制。
基于上述两个规则，我们设计了一种新的anchor对齐度量以显式度量anchor层面的任务对齐度。该对齐度量将集成到样本分配与损失函数中以动态提炼每个anchor的预测。
Anchor alignment metric 考虑到分类得分与IoU表征了预测质量，我们采用两者的高阶组合度量任务对齐度，公式定义如下：
其中，s与u分别表示分类得分与IoU值，而用于控制两者的影响。因此，在联合优化中起着非常重要的作用，它激励网络动态的聚焦于高质量的anchor。
Training sample assignment 正如已有研究表明，训练样例分配对于检测器的训练非常重要。为提升两个任务的对齐性，我们聚焦于任务对齐anchor，采用一种简单的分配规则选择训练样本：对每个实例，我们选择m个具有最大t值的anchor作为正样例，选择其余的anchor作为负样例。然后，通过新的损失函数(针对分类与定位的对齐而设计的损失函数)任务进行训练。
Classification objective 为显式提升对齐anchor的分类得分，同时降低不对齐者得分，我们在训练过程中采用替换正anchor的标签。然而，我们发现：当标签变小后网络无法收敛。因此，我们采用规范化进行替代，这里规范化后的满足以下两个属性：确保hard-instance的有效学习；保持instance之间的排序关系。
因此，我们采用一种简单的实例级规范化调整的尺度。因此BCE重写如下：
与此同时，我们采用Focal Loss缓解正负样本不平衡问题。因此，最终的分类损失定义如下：
Localization objective bbox(Bouding box)通过对齐的anchor(具有更大的分类得分、更精确的定位)预测得到，这样的bbox通常经过NMS后仍可以得以保留。此外，可以在训练阶段通过对损失加权选择高质量的bbox。因此，我们采用度量bbox的质量。类似分类目标，我们将定位优化目标调整为如下形式：
注：训练过程中，TAL的总损失则是上述两个损失相加。
上表对比了不同单阶段检测器的性能，从中可以看到：当采用ResNet101、ResNeXt-101-64x4d作为骨干时，TOOD分别取得了46.7AP与48.3AP指标，以3AP指标超过了ATSS， 以2AP指标超过了GFL；当采用ResNet101-DCN, ResNeXt-101-64x4d-DCN作为骨干时，TOOD取得了更大的性能提升(从48.3AP提升到51.1AP，提升幅度达2.8AP)。TOOD在单阶段目标检测方面取得了新的记录：51.1AP。
上表Head结构进行了消融分析，从中可以看到：当以“即插即用”方式将T-Head嵌入到不同单阶段检测器后均可以看到显著的性能提升，从0.7~1.9AP不等，同时具有更少的参数量与FLOPs。
上表在采样分配方面进行了消融分析，从中可以看到：相比此前最佳PAA，所提TAL+TAP组合可以取得1.6AP指标提升达到42.5AP。
上表对比了TOOD在Anchor-free与Anchhor-based两种检测器中的性能对比，可以看到：Anchor-free版TOOD与Anchor-based版TOOD具有相近的性能：42.5AP vs 42.4 AP。相比ATSS，TOOD取得了3.2AP指标提升。