TOOD: Task-aligned One-stage Object Detection

文章链接：https://arxiv.org/pdf/2108.07755.pdf
源码链接： https://github.com/fcjian/TOOD
单阶段目标检测通常通过优化目标分类和定位两个子任务来实现，使用两个并行分支的头部，这可能会导致两个任务之间的预测出现一定程度的空间错位。本文提出了一种任务对齐单阶段对象检测器(tood)，它以一种基于学习的方式显式地对齐两个任务。首先，设计了一种新型的任务对齐头(Task-aligned Head, T-Head)，它能更好地平衡学习任务交互和任务特定的特征，并通过任务对齐预测器来学习对齐。其次，提出任务对齐学习(TAL)，通过设计样本分配方案和任务对齐损失来明确地拉近(甚至统一)两个任务的最佳锚点。

一、文章简介：

目标检测旨在从自然图像中定位和识别感兴趣的目标，是计算机视觉中一个基础而又具有挑战性的任务。它通常被表述为联合优化对象分类和定位的多任务学习问题。分类任务的目的是学习集中在一个物体的关键或突出部分的鉴别特征。而定位任务则是精确定位整个物体及其边界。由于分类和定位学习机制的不同，两个任务学习到的特征的空间分布可能不同，使用两个单独的分支进行预测时，会造成一定程度的不对应。最近的单级物体探测器试图通过关注物体的中心来预测两个独立任务的一致输出。他们假设物体中心的一个锚(即，无锚检测器的锚点，或基于锚检测器的锚框)可能会对分类和定位给出更准确的预测。例如，FCOS和ATSS都使用中心度量分支来提高目标中心附近锚点预测的分类分数，对相应锚点的定位损失赋予更大的权重。FoveaBox将物体预定义中心区域内的锚点作为正样本。这种启发式设计取得了良好的效果，但这些方法可能存在两个局限性：
(1)分类和定位的独立性：最近的单级检测器通过使用两个独立的分支(即头部)并行地独立地进行对象分类和定位。这样的两个分支设计可能会导致两个任务之间缺乏交互，导致在执行它们时预测不一致。如下图的Result列所示，ATSS检测器识别出了Dining table的一个对象(用红色框表示的锚点)，但对Pizza的另一个对象(红色边框)的定位更为准确。

上图中上面一行为ATSS的检测结果(Result)和定位评分(IoU)的空间分布示意图，下一列是tood。真实值由黄色方框表示，白色箭头表示最佳锚点远离物体中心的主要方向。在Result列中，一个红/绿的小块是分类/定位最佳锚点的位置，而一个红/绿的方框是红/绿小块中锚点预测的一个物体边界框(如果它们重合，只显示红色的小块和方框)。
(2)任务不可知的样本分配：大多数无锚检测器使用基于几何的分配方案来选择目标中心附近的锚点进行分类和定位，而基于锚的检测器通常通过计算锚盒与真实值之间的iou来分配锚框。然而，用于分类和定位的最佳锚点往往是不一致的，并且可能会根据对象的形状和特征发生很大的变化。广泛使用的样本分配方案是任务不可知的，因此可能很难对两个任务做出准确而一致的预测，如上图中ATSS的Score和IoU分布所示。Result列还说明了最佳定位锚点(绿色框)的空间位置可能不在物体的中心，并且它与最佳分类锚点(红色框)没有很好地对齐。因此，在NonMaximum Suppression (NMS)期间，精确的包围框可能会被不那么精确的包围框抑制。
为了解决这些限制，TOOD通过设计一种新的头部结构，采用面向对齐的学习方法，更准确地对齐两个任务:
Task-aligned head：与传统的单阶段目标检测中使用两个分支并行实现分类和定位的头不同，本文设计了任务对齐头(Task-aligned head, T-head)来增强两个任务之间的交互。这使得两项任务能够更加协作，从而更准确地校准它们的预测。T-head在概念上很简单:它计算任务交互特征，并通过一种新的任务对齐预测器(TAP)进行预测。然后，它根据任务对齐学习提供的学习信号对齐两个预测的空间分布，如下所述。
(一)、学习任务对齐
为了进一步克服不对齐问题，作者提出了一种任务对齐学习(TAL)来明确地拉近两个任务的最佳锚点。通过设计一个样本分配方案和任务对齐损失来实现。样本分配通过计算每个锚点的任务对齐程度来收集训练样本(即正样本或负样本)，而任务对齐损失在训练过程中逐渐将最佳锚点统一起来，用于预测分类和定位。因此，在推理时，可以保留分类得分最高、联合定位最精确的边界框。

二、实现细节：

本文建议使用设计的任务对齐头(Task-aligned head, T-head)和新的任务对齐学习(Task-aligned Learning, TAL)来更明确地对齐两个任务。如下图所示，T-head和TAL可以协作来改进两项任务的一致性。具体来说，T-head首先对FPN特征进行分类和定位预测。然后，TAL根据一种新的任务对齐度量(测量两个预测之间的对齐程度)计算任务对齐信号。最后，T-head利用反向传播过程中从TAL计算的学习信号自动调整分类概率和定位预测。

上图展示的是TOOD的整体学习机制。首先，T-head对FPN特征进行预测。其次，预测用于计算每个定位点的任务对齐度量，TAL根据此度量为T-head生成学习信号。最后，T-head对分类和定位的分布进行相应的调整。其中，最对齐的锚通过概率图prob (probability map)获得更高的分类得分，通过学习偏移获得更准确的锚框预测。
Task-aligned Head
本文设计了一种高效的检测头结构来改进传统的一级探测器头部设计(如下图(a)所示)。在通过两个方面来实现:(1)增加两个任务之间的交互作用，(2)增强检测器学习对齐的能力。建议的T-Head如下图(b)所示，其中它有一个简单的特征提取器和两个任务对齐预测器(TAP)。

为了增强分类和定位之间的交互作用，使用特征提取器从多个卷积层学习大量任务交互特征，如上图(b)中的蓝色部分所示。该设计不仅方便了任务交互，而且为两个任务提供了多层次的特征和多尺度的有效接受域。形式上，$X_{fpn}$表示FPN特征，其中H、W和C分别表示高度、宽度和通道数。特征提取器使用连续的N层带有激活函数的转换层来计算任务交互特征。

其中$con{v}_k$和$\delta$分别为第k层和relu函数。因此，利用头部的单个分支从FPN特征中提取丰富的多尺度特征。然后，将计算得到的任务交互特征输入到两个TAP中进行分类和定位。
Task-aligned Predictor (TAP)
我们对计算得到的任务交互特征进行对象分类和定位，使两个任务能够很好地感知彼此的状态。但是，由于分支单一的设计，任务交互特征不可避免地会在两个不同的任务之间引入一定程度的特征冲突。对象分类和定位的任务有不同的目标，因此关注不同类型的特征(例如，不同的层次或接受域)。因此，提出了一种层关注机制，通过在动态计算这些特定于任务的特征来鼓励任务分解。如上图©所示，任务特定的特征是针对每个分类或定位任务单独计算的。
这里$w_k$是学习层注意力$w\in R^N$的第k个元素。W由跨层任务交互特性计算而来，能够捕获层与层之间的依赖关系:

这里的$f_{c_1}$和$f_{c_2}$指的是两个全连接层。σ是一个sigmoid函数，$x^{inter}$是$X_k^{inter}$的级联特征，他是通过对$X^{inter}$进行平均池化得到。最后，根据每个$X^{task}$预测分类或定位的结果.

其中$X^{task}$是$X_k^{task}$的级联特征，$con{v}_1是一个用来降维的1\ast 1卷积$，$Z^{task}$是用sigmoid转换的密集分类得分$P\in R^{H\ast W\ast 80}$,或者是一个回归框预测$B\in R^{H\ast W\ast 4}$
在预测步骤，可以通过调整两个预测P和B的空间分布，明确地对齐两个任务。而不像使用中心分支或IoU分支只能根据分类特征或定位来调整分类预测特征。通过使用计算任务交互特性来联合考虑两个任务，从而使这两个预测保持一致。在两个任务上分别执行对齐方法。使用空间概率图$M\in R^{H\ast W\ast 1}$来进行分类预测。

其中M是通过交互特征计算出来的，允许它学习两个任务在每个空间位置之间的某种程度的一致性。
同时，为了对定位预测进行对齐，进一步从交互特征中学习空间偏移图$O\in R^{H\ast W\ast 8}$，用于调整每个位置的预测锚框。学习到的空间偏移使最对齐的锚点识别周围的最佳边界预测。

其中索引(i, j, c)表示张量中第c个通道上的(i, j)个空间位置。上式采用双线性插值实现，由于b的通道维数很小，其计算开销可以忽略不计，值得注意的是，每个通道的偏移量都是独立学习的，这意味着物体的每个边界都有自己学习的偏移量。这样就可以更准确地预测四个边界，因为每个边界都可以从其附近最精确的定位点学习。因此，该方法不仅将两个任务对齐，而且通过为每一侧确定一个精确的定位点来提高定位精度。
下式为两个任务的表达式：

$conv1$和$conv3$都是用来降维的1*1卷积，利用提出的任务对齐学习(TAL)来完成M和O的学习。
Task Alignment Learning：我们进一步引入了任务对齐学习(TAL)，进一步指导我们的T-Head做出与任务对齐的预测。从任务对齐的角度，它根据设计的指标动态选择高质量的锚点。且同时考虑锚点分配和加权。它包括一个分配策略和新损失的函数，专门为调整两个任务而设计。
（一）、Task-aligned Sample Assignment
为了应对NMS，训练实例的锚点分配应满足以下规则:
(1)对齐良好的锚点应能联合精确定位预测出较高的分类分数;
(2)不对齐的锚应具有较低的分类评分，并随后予以抑制。基于这两个目标，设计了一个新的锚框对齐度量来明确地度量锚框级别上的任务对齐程度。该对准指标集成到样本分配和损失函数，以动态细化每个锚点的预测。
（二）、Anchor alignment metric
考虑到分类得分和预测边界框与真实值之间的IoU表示两个任务的预测质量，本文使用分类得分和IoU的高阶组合来衡量任务对齐程度。设计了以下度量来计算每个实例的锚框对齐程度：

其中s和u分别表示分类分数和IOU值。α和β被用来控制这两个任务的影响。值得注意的是，t在两个任务向任务对齐目标的联合优化中起着至关重要的作用。它鼓励网络从联合优化的角度动态关注高质量(即任务分配)锚点。
(三)、Training sample assignment
训练样本分配对于目标检测器的训练至关重要。为了提高两个任务的对齐性，应该关注与任务对齐的锚点，并采用一个简单的分配规则来选择训练样本:对于每个实例，我们选择t值最大的m个锚点作为正样本，而使用其余的锚点作为负样本。同样，训练是通过计算新的损失函数来完成的，这些损失函数是专门为分类和定位任务设计的。
Task-aligned Loss
（一）、Classification objective
为了显式地提高对齐锚点的分类分数，同时降低未对齐锚点的分类分数(即t)，在训练中使用t来代替正样本的二进制标签。然而，当正样本的标签(即t)随着α和β的增加而变小时，网络不能收敛。因此，需要使用一个归一化的t，即$\hat{t}$，来替换正样本的二进制标签，其中t由以下两个性质归一化:
(1)确保高质量样本的有效学习(通常所有对应的正锚点都有一个小的t）
(2)根据预测边界框的精度保持实例之间的顺序。因此，采用一个简单的实例级归一化来调整t的大小:t的最大值等于每个实例中最大的IoU值(u)。然后对分类任务的样本进行二值交叉熵(BCE)计算，可以重写为：

其中i表示对应于一个实例的正样本$N^{pos}$的第i个锚点。采用焦点损失进行分类，以减轻训练过程中负样本和正样本之间的不平衡。在正样本上计算的焦点损失可上式重新表示，分类任务的最终损失函数定义如下：

其中j表示负样本$N_{neg}$的第j个锚点，γ是聚焦参数.
(二)、Localization objective
由对齐良好的锚预测的边界框(即具有较大的t)通常具有较大的分类分数和精确的定位，并且这样的边界框在NMS期间更有可能被保留。此外，t可以用于选择高质量的包围框，更仔细地加权损失，以提高训练。从高质量的包围框中学习有利于模型的性能，而低质量的包围盒往往会产生大量信息较少、冗余的信号来更新模型，从而对训练产生负面影响。应用t值来测量边界框的质量。因此，通过关注对齐良好的锚点(t较大)来提高任务对齐和回归精度，同时在包围框回归期间减少错位锚点(t较小)的影响。与分类目标相似，对每个锚框基于$\hat{y}$计算的包围框回归损失重新加权，GIoU损失(LGIoU)可以重新表述如下:

其中b和$\hat{b}$表示预测的边界框和对应的真值框。TAL的总训练损失为$L_{c}ls$和$L_{r}eg$之和。

实验：

可以看出，实验结果还是很感人的