【论文阅读笔记】TOOD: Task-aligned One-stage Object Detection

论文代码：https://github.com/fcjian/TOOD

论文小结

  本文的目标，主要是解决不同任务之间预测的不对齐问题。任务不对齐，简单来说，就是分类的最佳anchor（point）和回归的最佳anchor（point）不一定是同一个，亦或者同个grid point的分类和回归所预测的实例对象不是同一个。

  为了解决这个问题，作者从模型结构的检测头或者训练策略中的标签分配中出发，分别设计了T-Head和TAL两个能相互结合的方案。其中TAL被后续文章和实践中较多使用。

  TAL的主要思路就是不事先赋予无法改变的标签，而是结合分类的预测 $s$ 和回归的预测生成的IoU分数 $u$来动态生成本次学习的标签。因此，TAL（Task Alignment Learning）中会有一个learning的说法。

  本文的对比实验和消融实验室未充分训练的对比，大致训练12个epoch来对比的。

论文简介

  由于分类和定位学习机制的差异，两个任务学习到的特征空间分布可能不同。当两个任务分支单独预测时，会导致一定程度的不对齐。如下图$1$所示，上面行是ATSS算法预测的结果，下面行是本文TOOD算法预测的结果。图中黄色框是餐桌的label，分类任务的最佳位置是红色色块，其边框预测为红框；回归任务的最佳位置是绿色色块。从图$1$可以容易看出，在原有方法训练管道得到的结果，可能会出现分类任务和定位任务的不对齐问题。

  在原先的方法（比如FCOS和ATSS）的认知中，对象的几何中心附近更容易预测对象（anchor-based在grid上，anchor-point），给予靠近中心的正样本分配更高的分数，更高的权重。这种启发式的方法取得很好的效果，但可能带来两个局限性：

1. 分类和定位的独立性。对象实例有重叠时，分类目标和定位目标不匹配；
2. 任务不可知标签分配。标签都是按照定位的几何中心分配的，但即使是定位问题，几何中心都不一定是最合适预测的。更何况标签的分配没有考虑分类任务的概率图（probably map）。分类和定位的最佳anchor通常是不一致的，并且依据目标的形状和特种有较大差异。因此很难对两个任务作出准确而一致的预测。

  为了解决分类任务和定位任务的不对齐问题，本文的贡献是提出了T-head和TAL两个方法。

  T-head，大致是为了增加两个任务之间的交互，以及增强检测器学习对齐的能力。总得来说，就是设计了一个新的检测头。

  TAL方法，目标是让两个任务的最佳anchor尽可能接近。其做法是设计一个简单的分配方案和任务对齐损失。样本分配通过每个anchor的任务对齐难度来收集训练的正负样本。而任务对其损失则在训练过程中逐渐统一预测分类和定位的最佳anchor。
  基本做法如下图所示，产生了cls和loc预测，再获得对齐指标，放入TAL中，进行标签分配和反向传播。

论文方法

  网络的整体设计也是按照“backbone - neck - head”结构来的。为简单和方便，TOOD在每个位置使用$1$个anchor（anchor point/box）。如上图$2$所示，T-Head和TAL可以协调地共同作用。原架构如下图$3(a)$所示，所提出的T-head如下图$3(b)$所示。

Task-aligned Head（T-Head）

  T-Head的目标是增加分类任务和定位任务的交互，以及增加两个任务之间的预测对齐能力。

  T-Head的架构图如上图$3(b)$所示，先通过多个卷积来堆叠任务交互特征，如公式（1）所示，然后再经过$2$个TAP（Task-aligned Predictor），用于对齐分类和定位。TAP的架构图如上图$3(c)$所示。
X k i n t e r = { δ ( c o n v k ( X f p n ) ) , k = 1 ，     ∀ k ∈ { 1 , 2 , . . . , N } δ ( c o n v k ( X k − 1 i n t e r ) ) ， k > 1 (1) X_k^{inter}=\begin{cases}\delta(conv_k(X^{fpn})),k=1， \ \ \ \ \forall k\in\{1,2,...,N\} \\ \delta(conv_k(X_{k-1}^{inter}))，k>1\end{cases}\tag{1} Xkinter​={δ(convk​(Xfpn)),k=1，    ∀k∈{1,2,...,N}δ(convk​(Xk−1inter​))，k>1​(1)

  由于TAP是一条单独的分支，任务交互特征不可避免地会在两个不同任务之间引入一定程度的特征冲突，所以作者提出来layer attention机制。通过在level级别动态计算任务特定特征，来鼓励任务分解，如公式（2）、（3），（4）。
X k t a s k = w ⋅ X k i n t e r , ∀ k ∈ { 1 , 2 , . . . , N } , (2) X_k^{task} = \mathcal{w}\cdot X_k^{inter}, \forall k\in \{1,2,...,N\}, \tag{2} Xktask​=w⋅Xkinter​,∀k∈{1,2,...,N},(2)
  下面公式中的$\delta$是全局池化层（GAP）；
$$w=\sigma (f{c}_2(\delta (f{c}_1(x^{inter}))))\text{\hspace{1em}(3)}$$
  下面公式中的$\delta$是ReLU非线性层；
$$Z^{task}=con{v}_2(\delta (con{v}_1(X^{task})))\text{\hspace{1em}(4)}$$

  如上图$3(c)$所示，对于预测的分类分数$P\in {\mathbb{R}}^{H\times W\times 80}$，从交互特征中，学习一个空间概率图（spatial probability map）$M\in {\mathbb{R}}^{H\times W\times 1}$来调整分类预测，完成对齐：$$P^{align}=\sqrt{P\times M}\text{\hspace{1em}(5)}$$

  同样的，对于预测的边框对象$B\in {\mathbb{R}}^{H\times W\times 4}$，从交互特征中学习一个空间偏差图（spatial offset maps）$\mathrm{O}\in {\mathbb{R}}^{H\times W\times 8}$来调整每个位置的bbox，完成对齐。$$B^{align}(i,j,c)=B(i+\mathrm{O}(i,j,2\times c),j+\mathrm{O}(i,j,2\times c+1),c)\text{\hspace{1em}(6)}$$

  $\mathrm{O}$和$M$从交互特征学习的公式如下：$$M=con{v}_2(\delta (conv1(X^{inter})))\text{\hspace{1em}(7)}$$$$\mathrm{O}=con{v}_4(\delta (conv3(X^{inter})))\text{\hspace{1em}(8)}$$

T-Head伪代码解释

  依据作者的描述是能够平衡学习两个任务交互特征，实际上思路大致是在neck之后先使用统一的管道来tune特征，就像早期的cls head与reg head未解耦时的状态。然后使用一些启发性的架构，比如channel attention来“解耦”cls特征和reg特征，channel attention的权重初步由全局池化层（GAP）得到，再通过卷积对该权重进行学习（类似于AsConv），作为attention的自适应权重。在得到分类和回归的推理结果上，作者也有一些启发性的设计。比如，分类任务首先预测一个分类分数，然后再预测一个分类概率相结合；回归任务首先预测一个bbox的xywh，再预测一个offset来进行调整。

feat_0 = feat_fpn
feat_{i} = leaky_relu(bn(Conv(feat_{i-1})))			                    # use for both cls task and reg task.
feat = torch.cat(feat_{i,...,n})weight = GAP(feat)
weight_cls = sigmoid(conv2(relu(conv1(weight)))				            # 从这儿解耦分类和回归，conv1/2/3/4表示不同卷积
weight_reg = sigmoid(conv4(relu(conv3(weight)))feat_cls = channel_attention(weight_cls, feat)
feat_reg = channel_attention(weight_reg, feat)# get cls prediction
cls_sigmoid = Conv3(256 -> num_anchors * 80)[feat_cls]                  # input is feat_cls
cls_prob = Conv1(256*6 -> 64)[feat] => relu => Conv3(64 -> 1)           # input is feat
cls_scores = \sqrt(cls_sigmoid.sigmoid() * cls_prob.sigmoid())# get reg prediction
reg_dist = Conv3(256 -> num_anchors * 4)[feat_reg].exp()                # input is feat_reg
reg_bbox = dis2bbox(reg_dist)
reg_offset = Conv1(256*6 -> 64)[feat] => relu => Conv3(64 -> 8)         # input is feat
bbox_pred = deform_Conv2d(reg_bbox, reg_offset)		                    # a library from mmlab/mmcv，add reg_offset for reg_bbox map. details see the formule.

Task Alignment Learning（TAL）

  在TAL中，主要有两点创新：

1. 动态选择最适合的anchor（在一个设计好的指标下）；
2. 它同时考虑anchor分配和权重，包括一个样本分配策略和专门为调整两个任务而设计的新Loss。

  anchor分配策略应该有如下$2$个规则：（1）一个对齐良好的anchor应该同时能预测一个高的分类分数和一个精确的定位；（2）没对齐的anchor应该有交底的分数，随后被NMS抑制。

Task-aligned Sample Assignment

  设定分类分数为$s$，预测和GT的IoU值为$u$，设计anchor级的对齐指标如下面公式，这样能同时衡量分类和定位的预测质量：$$t=s^{\alpha }\times u^{\beta }\text{\hspace{1em}(9)}$$

  对于每个实例，依据$t$值进行排序，选择最大值的$m$个anchors作为正样本，其他作为负样本。因此，每个训练迭代中，标签的分配都是动态赋予的。

Task-aligned Loss

  分类目标上，作者使用$t$值代替正样本anchor中的二元标签。
  但作者也发现，随着$\alpha$和$\beta$的上升，$t$值标签就会变小，进而导致无法收敛。所以作者将$t$进行正则化，标记为$\hat{t}$，来代替正样本anchor中的二元标签。$\hat{t}$的正则化应该遵循两项原则：（1）确保硬实例都能够高效地被学习（硬是难易样本的说法，比如Focal loss中的hard sample，此时所有的正anchor都只有很小的$t$值标签）；（2）根据预测框的精度来保持实例之间的等级。
  因此，作者采用了一个简单的实例级的正则化$\hat{t}$方法：$\hat{t}$的最大值等于每个实例的最大IoU值。（实际上就是换了个标签分配方案，可能为了方便写论文）。此时，经过BCE（Binary Cross Entropy）的损失函数为：$$L_{cls\_pos}=\sum _{i=1}^{N_{pos}}BCE(s_i,{\hat{t}}_i)\text{\hspace{1em}(10)}$$

  在分类任务上，使用focal loss，分类损失如下：其中$\gamma$是focal loss的超参数。$$L_{cls}=\sum _{i=1}^{N_{pos}}|{\hat{t}}_i-s_i{|}^{\gamma }BCE(s_i,{\hat{t}}_i)+\sum _{j=1}^{N_{neg}}{s}_j^{\gamma }BCE(s_j,0)\text{\hspace{1em}(11)}$$

  在定位任务上，作者通过关注对齐良好的anchors（具有较大的$t$值），来提高任务对齐和回归精度，同时减少边框回归过程中不对齐anchors的影响。

  基于$\hat{t}$来为每个anchor的bbox的loss重新分配权重。使用的损失函数为GIoU Loss，回归的损失公式如下：$$L_{reg}=\sum _{i=1}^{N_{pos}}{\hat{t}}_i{L}_{GIoU}(b_i,\overline{b_i})$$

论文实验

  COCO数据集，trainval135k为训练集，minival set为验证集。

  TOOD的backbone（ResNet-50，ResNet-101，ResNeXt-101-64x4d）都是在ImageNet上进行的预训练，检测头采用的是anchor-free的设计。本文实验中，anchor-based的组和anchor-free的性能相似。T-Head的连续卷积层数量$N=6$，Focal Loss的超参数$\gamma =2$。

消融实验

  消融实验中，使用的backbone为ResNet-50，默认训练周期为$12$个epochs（没有收敛完全）。

  下表1展示了T-head的消融实验结果，能涨AP$0.7-1.9$。

  TAL的消融实验如下表所示：

  Anchor-based和Anchor-free的消融实验如下：两者性能接近。

  $\alpha$和$\beta$的消融实验如下表所示：$\alpha =1.0$，$\beta =6$的情况下，得到的性能最好。

  TOOD与其他一阶目标检测算法的对比结果如下：其中每个模型使用不同分辨率（$480-800$）训练，训练$24$个epochs。

  TOOD整体结构的消融实验如下表所示：TOOD的设计 比 T-Head + ATSS 以及平行head + TAL的指标都好，表明T-head和TAL具有相辅相成的互补性。