TOOD Task-aligned One-stage Object Detection 论文学习

1. 解决了什么问题？

目标检测通过多任务学习的方式，协同优化目标的分类和定位。分类任务会学习目标的判别特征，关注于目标的显著性或关键区域，而定位任务则学习准确地定位目标的边界。因为定位和分类的学习机制不同，这俩任务学到的特征分布也不同。当这两个分支做预测时，会产生一定的错位现象。

• 分类和定位是独立的。两个独立的分支并行地做目标分类和定位，任务之间缺乏交流，会造成预测结果不一致。如下图红色格子，ATSS 识别的是餐桌，但定位的是披萨饼。
• Task-agnostic 样本分配。分类和定位的最佳 anchors 通常是不一致的，根据目标的形状和特性差异可能很大。常用的样本分配方法都是 task-agnostic，因此这俩任务很难作出准确而一致的预测。如下图中的绿色格子代表了最佳的定位 anchor，它不是目标的中心点，而且与分类的最佳 anchor（红色格子）没有对齐。这会造成 NMS 时，一个准确定位的边框被抑制掉。

如下图，上面一行是 ATSS 预测的分类得分和定位得分的空间分布，下面一行是 TOOD 预测的分类得分和定位得分的空间分布。黄框是 ground truth，红色格子是分类的最佳 anchor，绿色格子是定位的最佳 anchor，若二者重叠则只显示红色格子。红框和绿框分别是红色格子和绿色格子的 anchor 预测出的边框。白色箭头表示最佳 anchor 偏离目标中心点的主要方向。

2. 提出了什么方法？

TOOD 学习对齐这两个任务。作者首先设计了一个 task-aligned head (T-Head)，增强两个任务间的交流，在学习 task-interactive 和 task-specific 特征方面取得更好的平衡，更准确地对齐预测结果。其次，提出了 task alignment learning (TAL)，通过样本分配机制和 task-aligned loss 训练，拉近两个任务各自最优的 anchors 的距离。

2.1 Overview

TOOD 的整体流程采用了“主干-FPN-Head”的结构。TOOD 在每个位置使用一个 anchor，对于 anchor-free 检测器它就是一个 anchor point，对于 anchor-based 检测器它就是一个 anchor box。T-head 首先基于 FPN 特征预测分类和定位。然后 TAL 基于新提出的 task alignment metric 计算任务对齐信号，表示分类和定位预测的对齐程度。最后根据 TAL 反向传播的学习信号， T-head 自动地调节各类别的概率和边框位置。对得最齐的 anchor 会得到更高的分类得分，并学习偏移量使预测框更加准确。

2.2 Task-aligned Head

为了使 head 更加高效，作者从两个方面出发：

• 增加两个任务间的交流；
• 增强检测器学习对齐的能力。

T-head 如下图 (b) 所示，包括一个简单的特征提取器和两个 task-aligned predictors (TAP)。

为了增强分类和定位任务间的交流，使用一个特征提取器（多个卷积层）学习一组 task-interactive 特征，即上图(b) 蓝色部分。该设计不仅能提高任务间的交流，也能为这俩任务提供具有多尺度感受野的多层级特征。$X^{fpn}\in {\mathbb{R}}^{H\times W\times C}$代表 FPN 特征，其中$H,W,C$分别是 FPN 特征的高度、宽度和通道数。特征提取器使用$N$个卷积层和激活函数来计算 task-interactive 特征：
$$X_k^{inter}=\{\begin{array}{l}\delta ({\text{conv}}_k(X^{fpn})),k=1\\ \delta ({\text{conv}}_k(X_{k-1}^{inter})),k>1\end{array}\forall k\in \{1,2,...,N\}$$

其中${\text{conv}}_k$和$\delta$表示第$k$个卷积层和$\text{relu}$函数。在 head 里使用单分支结构提取丰富的多尺度特征，然后计算得到的 task-interactive 特征会输入进两个 TAP，学习对齐分类和定位。

2.2.1 TAP

在 task-interactive 特征上进行分类和定位，这两个任务能互相感知到对方的状态。但由于是单分支设计，task-interactive 特征难免会因任务不同而存在一定的特征冲突。作者提出了一个层注意力（layer attention）机制，动态地计算 task-specific 特征，使任务解耦。下图展示了 TAP，分别计算分类或定位的 task-specific 特征：

$$X_k^{task}={\mathbf{w}}_k\cdot X_k^{inter},\forall k\in \{1,2,...,N\}$$

${\mathbf{w}}_k$是学习得到的层注意力$\mathbf{w}\in {\mathbb{R}}^N$的第$k$个元素。$\mathbf{w}$是从 task-interactive 特征计算而来的，能够获取$X^{inter}$不同 layer 的依赖关系：

$$\mathbf{w}=\sigma (f{c}_2(\delta (f{c}_1({\mathbf{x}}^{inter}))))$$

$f{c}_1,f{c}_2$代表两个全连接层，$\sigma$是$\text{sigmoid}$函数，$\delta$是$\text{relu}$函数。将$X_k^{inter}$concat 起来得到$X^{inter}$，然后使用全局平均池化得到 ${\mathbf{x}}^{inter}$。最后，分类或定位的结果由每个 $X^{task}$预测得到：
$$Z^{task}={\text{conv}}_2(\delta ({\text{conv}}_1(X^{task})))$$

其中$X^{task}$是将$X_k^{task}$特征 concat 起来，${\text{conv}}_1$是$1\times 1$卷积，用于降维。$Z^{task}$然后使用$\text{sigmoid}$函数转换为分类得分$P\in {\mathbb{R}}^{H\times W\times 80}$，或者用$\text{distance-to-bbox}$转换为目标框$B\in {\mathbb{R}}^{H\times W\times 4}$。

2.2.2 Prediction alignment

预测时，通过调节两个预测（$P$和$B$）的空间分布，进一步对齐两个任务。以前的方法使用 center-ness 分支或 IoU 分支，基于类别特征或位置特征来调节类别预测。而本文则通过 task-interactive 特征综合考虑了两个任务，然后对齐这两个预测结果。如上图，作者使用一个空间概率图$M\in {\mathbb{R}}^{H\times W\times 1}$调节类别的预测：

$$P^{align}=\sqrt{P\times M}$$

$M$由 interactive 特征计算而来，学习每个空间位置上两个任务的对齐程度。

同时，为了对齐位置预测，从 interactive 特征学习一个空间偏移图$O\in {\mathbb{R}}^{H\times W\times 8}$，调节每个位置的预测框坐标。该偏移量使对得最齐的 anchor point 能识别到它附近最优的边界预测：

$$B^{align}(i,j,c)=B(i+O(i,j,2\times c),j+O(i,j,2\times c+1),c)$$

$(i,j,c)$表示张量中第$c$个通道的第$(i,j)$个位置。上式通过双线性插值实现，因为$B$的通道数不大，所以计算成本很低。注意，每个通道都会独立地学习偏移量，即目标的每条边都有一个偏移量。因为每条边都能学习它附近最准确的 anchor point，预测的四条边就能更加准确。因此，本文方法不仅能对齐定位和分类任务，也能通过识别每条边精确的 anchor point 来提升定位的准确率。

$M$和$O$从 interactive 特征中自动地学习：
$$M=\sigma ({\text{conv}}_2(\delta ({\text{conv}}_1(X^{inter}))))$$
$$O={\text{conv}}_4(\delta ({\text{conv}}_3(X^{inter})))$$

${\text{conv}}_1$和${\text{conv}}_3$是 2 个$1\times 1$的卷积层，用于通道降维。$M$和$O$通过 TAL 学习。T-head 是一个独立的模块，可以不需要 TAL。

2.3 Task Alignment Learning

使用 TAL 进一步指导 T-head 学习对齐分类和定位的预测。TAL 包括一个样本分配策略和一个新的损失函数。它从任务对齐的角度出发，动态地选取高质量 anchors，同时考虑 anchors 的分配和加权。

2.3.1 Task-aligned Sample Assignment

为了使用 NMS，anchor 分配应该满足以下两个条件：

• 对齐的 anchor 能同时预测出较高的类别得分和准确的边框位置；
• 错位的 anchor 的类别得分应该很低，会被抑制掉。

于是作者提出了一个新的 anchor 对齐度量，计算任务对齐的程度。在样本分配和损失函数中加入该度量，动态优化每个 anchor 的预测。

Anchor alignment metric. 类别得分和预测框与目标框间的 IoU 分别代表分类和定位任务的预测质量，作者将类别得分和 IoU 结合，表示俩任务的对齐程度。使用下式计算每个实例的各 anchor 的对齐度量：
$$t=s^{\alpha }\times u^{\beta }$$

$s$表示类别得分，$u$表示 IoU。$\alpha ,\beta$用于控制两项任务施加的影响。$t$在任务对齐优化的过程中扮演重要角色，使网络动态地关注于高质量（任务对齐）anchors。

Training sample assignment. 关注于任务对齐的 anchors，采用简单的分配规则来选取训练样本：对于每个实例，选取$m$个 $t$值最大 anchors 作为正样本，其余的 anchors 作为负样本，然后计算损失。

2.4 Task-aligned Loss

分类目标函数

为了抬高对齐的 anchors 的类别得分，降低那些没对齐的 anchors （$t$值偏小）的类别得分，训练时对$t$做归一化得到$\hat{t}$，代替 positive anchor 的二值标签。根据下面两个性质来归一化$\hat{t}$：

• 保证能有效地学习难例（即$t$值较小的 positive anchors）；
• 保留实例之间关于预测框精度的排序。

于是作者采用归一化方法来调节$\hat{t}$，对于每个实例， $\hat{t}$的最大值等于各 anchor 的 IoU$(u)$的最大值。对于分类任务，使用二元交叉熵来计算损失，将$\hat{t}$作为正样本的 ground-truth 标签，而非 1：

$$L_{cls\_pos}=\sum _{i=1}^{N_{pos}}BCE(s_i,{\hat{t}}_i)$$
$$BCE(s_i,{\hat{t}}_i)={\hat{t}}_i\cdot \log (s_i)+(1-{\hat{t}}_i)\cdot \log (1-s_i)$$

$i$表示某个实例的$N_{pos}$个 positive anchors 中的第$i$个 anchor。本文还用了 Focal Loss 缓解正负样本不均衡的问题。总的分类损失如下：
$$L_{cls}=\sum _{i=1}^{N_{pos}}|{\hat{t}}_i-s_i{|}^{\gamma }\cdot BCE(s_i,{\hat{t}}_i)+\sum _{j=1}^{N_{neg}}{s}_j^{\gamma }\cdot BCE(s_j,0)$$

$j$表示$N_{neg}$个 negative anchors 的第$j$个 anchor。$\gamma$和 Focal Loss 一文的含义相同，是调节系数。

定位目标函数

对齐的 anchors 预测出的边框通常置信度更高、边框更准确，才能在 NMS 时保留下来。训练时，$t$通过加权损失来提升高质量 anchor 的影响，降低低质量 anchor 的影响。高质量边框对模型有好处，而低质量边框则会产生大量冗余、无意义的信息。作者用$t$值计算边框的质量。利用$\hat{t}$对每个 anchor 的回归损失做加权，

$$L_{reg}=\sum _{i=1}^{N_{pos}}{\hat{t}}_i{L}_{GIoU}(b_i,{\overline{b}}_i)$$
$b,\overline{b}$分别是预测框和目标框。总的 TAL 训练损失是 $L_{reg}$和$L_{cls}$之和。