在过去几年中,YOLO 系列模型已成为实时目标检测领域的主流方法。许多研究通过修改架构、扩充数据以及设计新的损失函数,将基线水平提升到了更高层次。然而,我们发现尽管特征金字塔网络(FPN)和路径聚合网络(PANet)缓解了这一问题,但以往的模型仍存在信息融合方面的难题。
PART/1
概述
本研究提出了一种先进的聚合与分配(GD)机制,该机制通过卷积和自注意力操作实现。这种新设计的模型名为 Gold - YOLO,它增强了多尺度特征融合能力,并且在所有模型尺度下,都能在延迟与精度之间实现理想的平衡。此外,我们首次在 YOLO 系列模型中采用了 MAE 风格的预训练,这使得 YOLO 系列模型能够从无监督预训练中获益。Gold - YOLO - N 在 COCO val2017 数据集上达到了出色的 39.9% 的平均精度均值(AP),在 T4 GPU 上实现了 1030 帧每秒(FPS)的速度,与具有相似帧率的先前最优模型 YOLOv6 - 3.0 - N 相比,AP 提升了 2.4%。
PART/2
背景
目标检测作为一项基础的视觉任务,旨在识别物体的类别并定位其位置。它在众多应用领域有着广泛用途,如智能安防、自动驾驶、机器人导航以及医学诊断等。高性能且低延迟的目标检测器因可部署在边缘设备上,正受到越来越多的关注。
在过去几年里,研究人员对基于卷积神经网络(CNN)的检测网络展开了广泛研究,目标检测框架逐渐从两阶段(如 Faster RCNN和Mask RCNN)发展为单阶段(如YOLO),从基于锚框(如YOLOv3和YOLOv4)发展为无锚框(如 CenterNet、FCOS和YOLOX)。通过神经架构搜索(NAS)针对目标检测任务研究了最优网络结构,而其他文章则探索了通过知识蒸馏来提升模型性能的另一种途径。单阶段检测模型,尤其是 YOLO 系列模型,因其结构简单以及在速度与精度之间的平衡,在业界受到了广泛欢迎。
主干网络的改进也是视觉领域的一个重要研究方向。正如相关综述 [ A survey on vision transformer] 中所述,在精度与速度之间取得了平衡,而 [An image is worth 16x16 words: Transformers for image recognition at scale] 在精度方面表现出色。这些主干网络在从目标检测等高级任务到图像修复等低级任务的不同视觉任务中,提升了原始模型的性能。通过使用带有 Transformer 的编码器 - 解码器结构,研究人员构建了一系列类似DETR的目标检测模型,如DETR和DINO。这些模型能够捕捉物体之间的长距离依赖关系,使得基于 Transformer 的检测器能够达到与大多数经过精心优化的经典检测器相当甚至更优的性能。尽管基于Transformer的检测器表现出色,但与基于CNN的模型相比,其速度仍有不足。基于CNN的小尺度目标检测模型在速度 - 精度权衡方面仍占据主导地位,比如 YOLOX和YOLOv6-v3。
我们专注于实时目标检测模型,尤其是适用于移动设备部署的 YOLO 系列模型。主流的实时目标检测器由三部分组成:主干网络、颈部网络和头部网络。主干网络架构已得到广泛研究,而头部网络架构通常较为直接,由几个卷积层或全连接层构成。YOLO 系列中的颈部网络通常使用特征金字塔网络(FPN)及其变体来融合多层次特征。这些颈部模块基本遵循下图所示的架构。
然而,当前的信息融合方法存在一个显著缺陷:当需要跨层整合信息时(例如,融合第 1 层和第 3 层的信息),传统的类似 FPN 的结构无法无损地传输信息,这阻碍了 YOLO 系列模型实现更好的信息融合。
基于全局信息融合的理念,TopFormer在语义分割任务中取得了显著成果。在本文中,我们在 TopFormer 理论的基础上进行拓展,提出了一种新颖的聚合与分配机制(GD),通过全局融合多层次特征并将全局信息注入更高层次,以实现YOLO系列模型中高效的信息交换。这在不显著增加延迟的情况下,显著增强了颈部网络的信息融合能力,提升了模型在不同物体尺寸下的性能。具体而言,GD机制包含两个分支:一个浅层聚合与分配分支和一个深层聚合与分配分支,它们分别通过基于卷积的模块和基于注意力的模块来提取和融合特征信息。为了进一步促进信息流,我们引入了一个轻量级邻层融合模块,它在局部尺度上结合相邻层次的特征。我们的Gold - YOLO架构超越了现有的YOLO系列,有效地证明了我们所提方法的有效性。
为了进一步提高模型的精度,我们还引入了一种预训练方法,即使用MAE方法在 ImageNet 1K数据集上对主干网络进行预训练,这显著提高了模型的收敛速度和精度。例如,经过预训练的Gold-YOLO-S模型达到了46.4%的平均精度均值(AP),在相似速度下,超过了之前平均精度均值为45.0%的最优模型YOLOv6-3.0-S。
PART/3
新框架
如上图所示,YOLO 系列的颈部结构采用了传统的特征金字塔网络(FPN)结构,该结构包含多个分支用于多尺度特征融合。然而,它仅能完全融合相邻层级的特征,对于其他层级的信息,只能通过 “递归” 的方式间接获取。
在图中展示了传统 FPN 的信息融合结构:其中现有层级 1、层级 2 和层级 3 从上到下排列。FPN 用于不同层级之间的融合。当层级 1 从其他两个层级获取信息时,存在两种不同情况:
这种传递模式在计算过程中会导致大量信息丢失。层与层之间的信息交互只能交换中间层所选择的信息,未被选择的信息在传输过程中会被丢弃。这就导致了某一层级的信息只能充分辅助相邻层级,而对其他全局层级的辅助作用减弱。结果,信息融合的整体效果可能会受到限制。
为避免传统 FPN 结构在信息传输过程中的损失,我们摒弃了原有的递归方式,构建了一种新颖的聚合与分配机制(GD)。通过使用统一的模块收集并融合来自所有层级的信息,随后将其分配到不同层级,我们不仅避免了传统 FPN 结构固有的信息损失,还在不显著增加延迟的情况下增强了颈部的局部信息融合能力。
我们的方法因此能够更有效地利用主干网络提取的特征,并且可以轻松集成到任何现有的主干 - 颈部 - 头部结构中。
在我们的实现中,聚合与分配过程对应三个模块:特征对齐模块(FAM)、信息融合模块(IFM)以及信息注入模块(Inject)。
为增强模型对不同尺寸物体的检测能力,我们开发了两个分支:低阶段聚合与分配分支(Low - GD)和高阶段聚合与分配分支(High - GD)。这两个分支分别提取并融合大尺寸和小尺寸的特征图。
如上图所示,颈部的输入由主干网络提取的特征图 B2、B3、B4、B5 组成,其中 Bi ∈ RN×CBi×RBi 。这里 N 表示批量大小,C 表示通道数,R = H × W 表示维度。此外,RB2、RB3、RB4 和 RB5 的维度分别为 R、\(\frac{1}{2}R\)、\(\frac{1}{4}R\) 和 \(\frac{1}{8}R\)。
在这个分支中,我们选取主干网络输出的 B2、B3、B4、B5 特征进行融合,以获取保留小目标信息的高分辨率特征。其结构如下图(a)所示。
在低阶段特征对齐模块(Low - FAM)中,我们采用平均池化(AvgPool)操作对输入特征进行下采样,以实现尺寸统一。通过将特征调整为该组中最小的特征尺寸(RB4 = \(\frac{1}{4}R\)),我们得到 Falign。Low - FAM 技术确保了信息的高效聚合,同时通过 Transformer 模块将后续处理的计算复杂度降至最低。
目标对齐尺寸的选择基于两个相互矛盾的考量:(1)为保留更多低层次信息,较大的特征尺寸更为可取;然而,(2)随着特征尺寸增加,后续模块的计算延迟也会上升。为控制颈部部分的延迟,有必要维持较小的特征尺寸。
因此,我们选择 RB4 作为特征对齐的目标尺寸,以在速度和精度之间达成平衡。
低阶段信息融合模块(Low - IFM)的设计包含多层重参数化卷积块(RepBlock)和一个拆分操作。具体而言,RepBlock 将 Falign(通道数 = CB2 + CB3 + CB4 + CB5 之和)作为输入,并生成 Ffuse(通道数 = CB4 + CB5)。中间通道数是一个可调整的值(例如 256),以适配不同的模型规模。RepBlock 生成的特征随后在通道维度上被拆分为 Finj_P3 和 Finj_P4,然后分别与不同层级的特征进行融合。
如上图(b)所示,高阶段聚合与分配分支(High - GD)融合由低阶段聚合与分配分支(Low - GD)生成的 {P3, P4, P5} 特征 。
高阶段特征对齐模块(High - FAM)由平均池化(avgpool)组成,用于将输入特征的维度降低到统一大小。具体来说,当输入特征的尺寸为 {RP3, RP4, RP5} 时,平均池化将特征尺寸缩小到该组特征中的最小尺寸(RP5 = \(\frac{1}{8}R\))。由于 Transformer 模块提取的是高级信息,池化操作有助于信息聚合,同时降低了 Transformer 模块后续步骤的计算需求。
高阶段信息融合模块(High - IFM)包括 Transformer 模块(下面将详细解释)和一个拆分操作,该过程分为三步:
为了更有效地将全局信息注入到不同层级,我们从图像分割的经验中获得启发,采用注意力操作来融合信息,如上图所示。具体而言,我们同时输入局部信息(即当前层级的特征)和全局注入信息(由信息融合模块 IFM 生成),分别记为 Flocal 和 Finj。我们对 Finj 使用两个不同的卷积层进行计算,得到 Fglobal_embed 和 Fact。同时,使用卷积层对 Flocal 进行计算得到 Flocal_embed。然后,通过注意力机制计算得到融合后的特征 Fout。由于 Flocal 和 Fglobal 的尺寸不同,我们采用平均池化或双线性插值,根据 Finj 的尺寸对 Fglobal_embed 和 Fact 进行缩放,以确保正确对齐。在每次注意力融合结束时,我们添加重参数化卷积块(RepBlock),以进一步提取和融合信息。
PART/4
实验及可视化
在 COCO2017 验证集上与其他 YOLO 系列检测器的比较。帧率(FPS)和延迟是在相同环境下,使用 TensorRT 7,基于 Tesla T4 以 FP16 精度进行测量的。我们所有的模型都训练了 300 个 epoch。我们模型的精度和速度性能均在输入分辨率为 640×640 的情况下进行评估。“†” 表示使用了自蒸馏方法,“⋆” 表示使用了基于掩码图像建模(MIM)的预训练方法。
上图中展示的 YOLOv5、YOLOv6、YOLOv7、YOLOv8 以及我们的 Gold - YOLO 颈部的类激活映射(CAM)可视化结果。可以观察到,我们的模型对目标的检测区域赋予了更高的权重。
并且,我们在下图中对比了 Gold - YOLO 和 YOLOv6 的颈部 CAM 可视化结果。