目标检测即插即用 | Uni-Head这个学霸让ATSS/TOOD/RetinaNet/FreeAnchor通通涨点

点击下方卡片，关注「AI视界引擎」公众号

检测Head是目标检测器中的一个关键组件，负责执行分类和定位功能。然而，常用的并行检测Head往往缺乏全方位的感知能力，如Deformation感知、全局感知和跨任务感知。尽管有许多方法试图从单一方面增强这些能力，但实现全面统一的解决方案仍然是一个重大挑战。为了应对这一挑战，作者开发了一种创新的检测Head，称为UniHead，可以同时统一三种感知能力。更具体地说，作者的方法包含：

1. 引入了Deformation感知，使模型能够自适应地采样目标特征；
2. 提出了双轴聚合Transformer（DAT），以熟练地建模长距离依赖关系，从而实现全局感知；
3. 设计了一个跨任务交互Transformer（CIT），促进了分类和定位分支之间的交互，从而使这两个任务保持一致。

作为一种即插即用的方法，提出的UniHead可以方便地集成到现有的检测器中。在COCO数据集上进行的大量实验证明，作者的UniHead可以显著改善许多检测器的性能。例如，UniHead可以在RetinaNet中获得+2.7的AP增益，在FreeAnchor中获得+2.9的AP增益，在GFL中获得+2.1的AP增益。代码将会公开发布。

1、简介

目标检测是一项基础但具有挑战性的任务，旨在在图像中定位和识别感兴趣的目标。近年来，目标检测取得了显著的进展，包括Backbone网络设计、特征融合网络优化和有效的训练策略。尽管取得了令人印象深刻的突破，但目前的研究尚未充分探讨目标检测中检测Head的关键作用。

检测Head的主要目的是利用从Backbone网络提取的复杂特征执行准确的分类和定位任务。一种普遍的方法是并行Head，它通过使用堆叠卷积来建立两个独立的分支，以学习与任务相关的特征，如图1(a)所示。然而，尽管广泛使用，但这种方法存在一些局限性，尤其是对于理想的检测器而言，缺乏必要的感知能力。

具体来说，理想的检测器应具备三种关键的感知能力：Deformation感知、全局感知和跨任务感知。具体而言，自然目标表现出两种类型的多样性：几何Deformation多样性、尺度和形状多样性。前者要求检测器能够自适应地采样与目标目标的几何Deformation相对应的特征（Deformation感知），而后者要求具有全局感知能力，以捕捉全局特征并建模长距离依赖关系。

此外，许多研究发现，分类预测和定位预测可能不一致，即具有高分类分数的框并不总是能够准确定位。因此，检测Head需要跨任务感知，以全面整合来自两个任务的监督信息，从而实现对分类和定位的一致检测。

然而，常用的并行Head由于卷积属性和并行结构的限制，无法提供这些感知能力。特别是，卷积机制在Deformation感知和全局感知方面受到限制，因为它在固定的局部位置采样输入特征图。并行结构独立执行两个任务，并且缺乏跨任务感知以执行它们之间的交互，如图1(a)所示。

最近，研究人员试图通过提出新颖的检测Head来减轻上述问题。例如，Dynamic Head集成了多种注意力机制，统一了尺度感知和空间感知，从而缓解了卷积机制的局部性偏差。GFL和TOOD结合了来自两个任务的监督信息（例如分类分数和IOU），以提高模型的跨任务感知能力，从而实现一致的检测预测。然而，现有方法倾向于增强三个感知方面中的一个，缺乏一个可以同时增强所有这些感知的统一方法。

因此，在本文中，作者提出了一种即插即用的检测Head，称为UniHead，以从综合和整合的角度解决上述问题。UniHead包括不同的模块，以赋予检测器全面的感知能力，同时不会增加额外的计算或架构负担。UniHead的总体框架如图1(b)所示。

具体而言，作者首先通过经典的DCN向UniHead引入了Deformation感知。这使模型能够在远离局部和固定形状的空间位置进行采样，从而更有效地感知Deformation的目标。其次，作者提出了一种双轴聚合Transformer（DAT），以有效建模长距离依赖关系。DAT在通道压缩空间中水平和垂直轴上并行执行自注意力，这个设计使其能够在一个单一模块中捕捉全局信息，同时保持低计算复杂性。更重要的是，作者引入了一个跨任务交互Transformer（CIT），它使用交叉注意力来促进分类和定位任务之间的交互。通过整合CIT，UniHead不仅可以捕获一个任务的丰富上下文，还可以利用来自另一个任务的相关上下文，从而明确促进两个任务之间的相互对齐。

基于上述精心设计的三个模块，UniHead可以有效地增强多个感知能力。作者的UniHead可以方便地作为插件轻松集成到现有的检测器中。作者在MS-COCO基准上进行了广泛的实验，以展示作者方法的有效性。

如图2所示，UniHead显著提高了经典检测器的性能。具体而言，将UniHead应用于以ResNet-50为Backbone的RetinaNet时，实现了39.2的平均精度（AP），显著提高了RetinaNet Baseline 2.7的AP。使用在ImageNet-22K数据集上预训练的强大Backbone Swin-B，作者的UniHead达到了54.3的AP，展示了作者方法的潜力和与大型模型的兼容性。

本工作的主要贡献可以总结如下：

• 作者提出了一种新颖的检测Head，即UniHead，通过共同增强检测器的三种重要感知能力，即Deformation感知、全局感知和跨任务感知，来提高检测性能。

• 作者设计了一个双轴聚合Transformer（DAT）来有效而高效地捕捉全局特征。此外，作者引入了一个精心设计的跨任务交互Transformer（CIT），以促进分类和定位任务之间的交互。

• 作者在MS-COCO基准上验证了装备作者方法的各种检测器，结果表明，作者的方法可以在更少的计算成本下持续提高这些检测器的性能，提高了1.7 ∼ 2.9的AP。

2、相关工作

A. 目标检测

近年来，在目标检测领域见证了蓬勃发展，主要有两类目标检测器：两阶段和一阶段检测器。

两阶段检测器，如R-CNN系列，在第一阶段使用区域建议网络（RPN）生成区域Proposals，然后在第二阶段优化这些Proposals的预测。

与两阶段范式不同，一阶段检测器消除了区域Proposals步骤，而是直接对边界框进行分类和回归。然而，早期的一阶段检测器在检测性能方面落后，直到RetinaNet的出现，它使用Focal Loss来解决类别不平衡问题。在RetinaNet之后，各种检测器消除了广泛使用的Anchor，发展出Anchor-Free检测器，它们使用中心和角点来表示目标。一些研究人员还提出了新颖的损失和训练策略来提高检测器的性能。

B. 检测Head

检测Head是检测器的关键组成部分，广泛采用的并行Head是默认选择。然而，最近的研究揭示了这种标准检测Head在实现最佳检测性能方面存在不足。例如，Double-Head指出，全连接Head更适合分类任务，而卷积Head更适合定位任务。GFL通过引入软标签丰富了检测Head的概念。它利用了定位任务的IOU分数作为分类标签，形成了一个联合表示定位质量和分类的表示。TOOD设计了一个任务对齐Head，旨在增强分类和定位任务之间的互动。Dynamic Head使用注意力机制增强了检测Head关于尺度和空间性的感知能力。

尽管这些值得赞扬的进展，但当前的方法倾向于分别解决特定的子问题。本文介绍了一个综合性的检测Head，从多个角度增强了模型的检测能力。

C. 注意力机制

注意力机制在人类感知中起着至关重要的作用。受此启发，它们也广泛应用于深度学习中，以进一步提高模型的性能。在各种注意力机制中，DCN可以被视为一种特殊的注意力机制，它在普通卷积基础上添加可学习的偏移，以从局部区域采样空间位置。

近年来，在计算机视觉领域，Transformer因其在建模长距离依赖关系方面的优势而备受赞誉。尽管Transformer可以为模型提供全局感受野，但其复杂性与图像分辨率的二次关系限制了其在高分辨率下游任务（包括目标检测）中的应用。

因此，已经提出了大量算法来改善这个问题，包括引入全局Token，减小注意力的空间尺寸，设计新颖的注意力机制（例如，局部窗口注意力，轴向注意力，十字形注意力，交叉窗口注意力）。

在本文中，作者利用现有的注意力机制或设计新的注意力机制来增强感知能力。为了最小化计算开销，作者还采用了轻量级设计，使作者的UniHead既有效又高效。

3. UniHead

广泛采用的并行Head未能有效处理检测任务的两个特性：自然环境中的物体多样性以及分类和定位之间的相互作用，这是由于卷积和并行结构的限制所致。

为了改善这一问题，作者提出了一种新颖的检测Head，称为UniHead，如图1(b)所示。

UniHead可以以统一的形式为模型提供三种能力：Deformation感知、全局感知和跨任务感知。

• Deformation感知：Deformation感知使模型能够自适应地学习与目标相关的特征，而不是被困在固定的窗口和局部位置。

• 全局感知：全局感知允许模型感知全局特征并建模长距离依赖关系，从而更准确地检测具有不同尺度的目标。

• 跨任务感知：跨任务感知可以在两个任务之间执行交互，为每个任务引入额外的监督信息，从而实现更一致的预测。

作者利用或引入专门设计的模块来实现这些感知，并将它们统一起来构建作者的UniHead。通常情况下，给定多尺度特征图，UniHead执行这三个模块来捕获三种感知特征。然后，输出特征用于执行分类和定位。通过这个框架，UniHead试图将三种感知能力统一到一个Head部中。

A. Deformation感知

自然场景中的目标复杂，具有各种内容和几何变换。使用固定 Kernel （例如3×3）的普通卷积难以有效处理这种具有挑战性的情况。

因此，作者通过使用DCN将Deformation感知（DP）引入到UniHead中。DCN可以通过在多尺度特征上学习的偏移和尺度来感知目标的变换。对于给定的3×3 Kernel 和偏移，位置处的Deformation学习过程可以表示为：

其中表示一个3×3卷积操作的采样位置数量，表示第个采样位置的卷积权重。表示特征上的双线性插值。和分别表示第个采样位置的预测偏移和尺度。

预测偏移使特征采样不受限于固定位置，而调制尺度调节每个采样位置的重要性。通过这些自适应采样偏移和调制尺度，作者将可Deformation表示能力引入到模型中，从而有助于检测具有复杂形状和多样外观的目标。

B. 全局感知

图像中存在各种尺度和形状不同的目标，这要求检测器捕捉全局特征以定位它们的复杂边界。为了解决这个问题，作者提出了一个双轴聚合Transformer（DAT）来建模长距离依赖性，从而增强检测器的全局感知能力。DAT包括两个部分：高效的双轴注意和跨轴聚合块。

高效的双轴注意

尽管Transformer已经证明了其在全局建模能力方面的有效性，但其巨大的计算负担限制了其更广泛的应用。替代方法，如局部窗口注意力和空间降维操作，通常要么舍弃全局能力，要么丢失一些空间信息。为了克服这些限制，作者提出了高效的双轴注意力（EDA）。

如图3(a)所示，EDA并行使用水平和垂直轴向注意来建模长距离依赖关系。对于水平轴向注意力，输入特征 被均匀分成H个不重叠的水平轴向Stripe，每个带有W个Token。设表示第个Stripe，的自注意力可以表示为：

其中 , , 分别表示输入上的Query、Key和Value的投影矩阵。请注意，在计算Query、Key和Value时，使用了通道缩减操作，使得注意力在通道压缩空间中执行。和表示共享的值和投影矩阵，并且将在水平和垂直轴向注意中使用。是的维度。表示的水平注意力输出。垂直轴向自注意力可以类似地推导出，其输出表示为。最后，两个部分的输出沿通道维度串联。该过程表示为：

其中Cat表示按通道串联；是用于特征融合的常用投影矩阵。EDA的复杂性为：

因此，EDA将注意力的复杂性降低到与图像高度或宽度的二次关系（或），而不是与图像分辨率的二次关系（）。

与先前的轴向Stripe自注意力相比，作者的EDA是在通道压缩空间中计算的，并共享相同的值映射。这两个操作使得作者的EDA能够更高效地捕获全局特征。

跨轴聚合块

尽管EDA可以有效地建模Token之间的长距离依赖关系，但由于自注意力缺乏感知的归纳偏差，它在学习局部信息方面表现不佳。此外，轴向注意不能直接聚合跨轴信息。为了补充EDA的局部性并实现全局和局部的耦合，作者提出了一个跨轴聚合块（CAB）。如图3(b)所示，CAB应用于值（V）映射上，并与轴向注意并行运行。从公式上看，这个过程可以表示为：

其中和分别表示水平轴向注意的值映射和输出。CAB的功能可以通过一个3×3的深度卷积轻松实现。CAB提供了一种更灵活的机制，不仅提供了具有位置信息的注意力，还使不同轴向Stripe的注意力之间能够进行交互和聚合，如图3(c)所示。

有了作者提出的EDA和CAB，作者的双轴聚合Transformer（DAT）可以有效地建模长距离依赖关系，实现全局感知，从而更精确地定位。

C. 跨任务感知

目标检测是分类和定位的综合。检测Head需要利用两个任务的信息来进行一致的预测，而不是独立执行这两个任务。换句话说，检测器需要输出具有精确位置和高分类置信度的边界框。因此，作者提出了一个跨任务交互Transformer（CIT），以弥补模型执行跨任务交互的能力。

如图4所示，CIT具有两个重要组件：跨任务通道注意和局部增强块。

跨任务通道注意

作者利用通道级别的跨任务注意力来执行分类和定位任务之间的交互。通道级别的注意力可以从两个方面使模型受益：

1. 将注意力的复杂性降低为与图像大小呈线性关系；
2. 增强模型的通道级全局感知，同时与侧重于全局空间感知的DAT相互补充。跨任务注意力使模型能够利用一个任务的特征来通知和引导另一个任务的特征学习过程。

具体来说，在将特征输入CIT之前，作者首先使用条件位置编码为分类和定位特征编码位置信息，输出分别表示为和。如图4的左侧所示，当将CIT应用于分类分支时，跨任务通道级注意力（CCA）可以描述为：

其中、、和 分别是输入和的Key和Value的投影矩阵； 表示输入的Query的投影矩阵；CA表示计算通道之间的注意力的通道级注意力。

如公式6所示，作者使用中的跨任务监督信息作为Query来引导分类特征的表示学习。此外，作者将两个任务的特征连接起来以生成Key和Value，从而进一步融合跨任务信息。

局部增强块

由于通道级别的注意力缺乏学习局部空间特征的能力，作者提出了一个局部增强块（LEB）来缓解这个问题。LEB轻量而有效，建立在具有不同 Kernel 大小的深度卷积之上。从公式上看，它可以描述为：

其中DW表示深度卷积；表示跨任务通道级注意力（CCA）输出的特征。两个1×1的深度卷积可以进一步调节每个通道的重要性，它们是通道级别的尺度。3×3的深度卷积可以学习局部信息，并为模型提供归纳偏差的补偿。

4、实验

1. 主要结果

应用于经典检测器

作者通过将UniHead插入到流行的目标检测器中，包括RetinaNet、FCOS、RepPoints、FreeAnchor、ATSS和GFL，来评估UniHead的有效性和泛化能力。这些选择的检测器代表了广泛的目标检测框架，包括基于Anchor、基于中心的Anchor-Free、基于关键点的Anchor-Free和强 Baseline 。

为了公平比较，作者使用感知模块数配置为1、2、2，以保持与并行Head相当的复杂性。

结果报告在表I中。请注意，RepPoints的并行Head中的卷积层数比其他选定的检测器少，因此作者的方法的复杂性略高。如表I所示，UniHead始终显著提高所有检测器的性能，例如在RetinaNet上提高了2.7个AP，在FreeAnchor上提高了2.9个AP，即使参数和计算量更少。这些结果说明了作者方法的有效性和效率。

与不同的 Backbone 网络合作

作者展示了作者的方法与各种 Backbone 网络的兼容性，并与最先进的（SOTA）检测器进行比较。作者利用GFLv2作为作者的检测框架，感知模块的数量为1、3、3。作者对基于CNN的 Backbone 网络进行24个Epoch的训练，对基于Transformer的 Backbone 网络进行36个Epoch的训练。作者在训练过程中采用多尺度训练策略，并在单一尺度上测试模型。

为了公平比较，作者报告了其他模型使用多尺度训练和单一尺度测试的结果。如表II所示，作者的UniHead在强大的 Backbone 网络下取得了更加令人印象深刻的性能，远远超过了对应的SOTA模型。

具体来说，与最佳检测器TOOD相比，使用相同设置，作者的方法在ResNet-101 Backbone 网络上优于其1.0个AP。与先进的DETR系列相比，作者的UniHead可以在更少的训练Epoch（24个Epoch对比50个Epoch）内实现更高的准确性。作者的UniHead还与基于Transformer的 Backbone 网络（Swin-T）兼容，并且在0.9个AP上优于出色的级联Mask R-CNN。

此外，作者使用了在ImageNet- 22K数据集上预训练的Swin-B作为 Backbone 网络，以探索大型模型的性能改进。如表II所示，使用更强大的 Backbone 网络和预训练模型，作者的UniHead将性能显着提高到54.3个AP，显示了作者方法进一步改进的潜力。

B. 消融研究

作者进行了详细和全面的消融研究，以展示作者的UniHead的有效性和效率。对于消融研究，作者将UniHead应用于具有ResNet-50 Backbone 网络的RetinaNet，并在没有多尺度训练的情况下进行了12个Epoch的训练。

感知模块的有效性

为了评估每个提出的模块的有效性，作者进行了一项对照研究，其中作者消除了RetinaNet中使用的并行Head，并逐渐添加不同的感知模块。

研究的结果见表III。作者的发现表明，当仅应用一个感知模块到RetinaNet时，它无法获得令人满意的结果。有趣的是，作者观察到仅使用跨任务感知就能胜过 Baseline ，强调了跨任务交互的重要性。

随着更多感知模块的添加，作者观察到RetinaNet的性能显著提升。例如，当同时使用全局感知和跨任务感知时，作者的UniHead即使在计算复杂性更低的情况下也能优于 Baseline 1.4个AP。最后，作者的完整UniHead集成了三种感知能力，可以显著提高 Baseline 2.7个AP。

感知模块的数量

作者通过堆叠不同数量的感知模块来评估作者方法的效率和可扩展性。Deformation感知模块可以被视为隐式的位置编码，因此作者将其数量保持为1。

如表IV所示，作者的方法可以通过堆叠更多的模块来持续获得显著的性能提升，直到模块数量达到1、4、4。值得注意的是，即使在计算复杂性更低的情况下，作者的方法（每种感知类型使用一个模块）仍然能够显著优于 Baseline 1.7个AP。这表明了作者方法的效率。

此外，当作者采用模型配置1、3、3时，UniHead实现了39.8个AP，显著提高了 Baseline 3.3个AP，说明了作者方法的强大和可扩展性。

Deformation感知模块的有效性

作者通过将提出的Deformation感知模块（DPM）替换为标准的3×3卷积来评估Deformation感知的有效性。如表V所示，DPM的性能优于普通卷积0.4个AP。这表明，DPM引入的Deformation感知可以增强模型的表示能力和检测性能。

与其他Transformer模块的比较

作者将提出的双轴聚合Transformer（DAT）与其他经典的Transformer模块进行比较，包括Swin-Transformer块、CSwin-Transformer块和轴向注意力。

为了公平比较，作者将DAT替换为这些Transformer模块，同时保持其他结构不变。结果报告在表VI中。显然，作者提出的DAT在效率和性能方面优于其他Transformer模块。

条纹宽度

表VII呈现了全局感知模块中条纹宽度的消融分析。结果表明，扩展注意力区域并不会导致模型性能进一步提高；实际上，它甚至可能导致性能下降。这表明，DAT的潜力来自于其高效的并行结构，使模型能够感知全局信息。

跨轴聚合块的有效性

在表VIII中，作者呈现了在作者方法中使用跨轴聚合块（CAB）的结果。CAB有效地整合了水平和垂直轴向注意力的局部和跨轴信息，从而获得性能改进。

跨任务感知的有效性

为了评估跨任务感知的效果，作者将提出的跨任务通道注意力（CCA）替换为通道内自注意（CSA），结果见表IX。结果表明，作者提出的CCA在不添加额外参数和计算的情况下优于CSA 0.3 AP，这表明整合跨任务信息可以有效地指导特征学习，从而实现更高的性能。

局部增强块的有效性

作者消融跨任务感知模块中的局部增强块（LEB）。如表X所示，所提出的LEB是轻量级的，能够提高模型0.2 AP，证明了局部增强的重要性。

E. 可视化

作者使用RetinaNet进行检测，并在将其并行Head替换为UniHead后进行检测结果的可视化。

如图5所示，UniHead可以更有效地检测复杂场景中的目标，其中目标具有各种尺度和多样的几何变换。值得注意的是，图5的第三列突出显示了UniHead在有效检测小目标方面的能力。这种出色的检测性能，特别是对于形状复杂和小目标，突显了Deformation感知和全局感知的有效性。

此外，如图5的最后一列所示，作者的方法可以提供具有更高分类分数的精确检测框，表明UniHead可以帮助检测器在分类和定位方面输出更一致的结果。

F. 进一步的泛化验证

作者进一步将UniHead泛化到两个场景中的代表性两阶段模型，包括使用Faster R-CNN进行目标检测和使用Mask RCNN进行实例分割。

尽管两阶段模型没有并行Head，但作者的UniHead可以轻松集成到它们中。作者使用UniHead替换了Region Proposal Network（RPN）中的建议预测之前的卷积。

请注意，RPN中的卷积数量为1，因此为了避免增加复杂性，UniHead在通道缩减后使用。

如表XI所示，在这些场景中，作者的模型仍然取得了非常令人印象深刻的性能。具体来说，UniHead将Faster R-CNN提高了1.6个AP，将Mask RCNN提高了1.1个AP，充分验证了作者方法的泛化能力。

5. 参考

[1]. UniHead: Unifying Multi-Perception for Detection Heads.

6. 推荐阅读

港科大重组YOLOv5模块 | HIC-YOLOv5优化小目标问题，虽不新颖但很有用!

可以取代NMS？IOU感知校准，取代NMS，成就YOLOX等目标检测更快的后处理方法

DMKD蒸馏 | 模型学习，空间/通道信息我都要！超越FKD、FGD、MGD以及AMD方法！

点击上方卡片，关注「AI视界引擎」公众号