LW-DETR: 用轻量Transform来代替YOLO的实时目标检测

ronghuaiyang

于 2024-09-10 12:00:37 发布

阅读量385

点赞数

文章标签： YOLO 目标检测人工智能计算机视觉

原文链接：https://mp.weixin.qq.com/s?__biz=Mzg5ODAzMTkyMg==&mid=2247497747&idx=1&sn=f0d3f9bc06038a5514727a23735b6dd5&chksm=c1cab0eadd6f72cea8c583804a1bcd0fefce1cefd85058260d6acaeda04c4bf6a48573e39b04&scene=126&sessionid=0

版权

导读

在本文中，我们介绍了一种轻量级检测Transformer（LW-DETR），其在实时目标检测方面优于YOLO系列模型。该架构简单地堆叠了一个ViT编码器、一个投影器和一个浅层DETR解码器。

摘要

在本文中，我们介绍了一种轻量级检测Transformer（LW-DETR），其在实时目标检测方面优于YOLO系列模型。该架构简单地堆叠了一个ViT编码器、一个投影器和一个浅层DETR解码器。我们的方法利用了最近的先进技术，例如有效的训练技术，包括改进的损失函数和预训练，以及交错窗口和全局注意力机制来降低ViT编码器的复杂度。我们通过聚合多层特征图改进了ViT编码器，并在ViT编码器中整合了中间和最终特征图，从而形成更丰富的特征图，并引入了窗口主特征图组织方式以提高交错注意力计算的效率。实验结果表明，所提出的方法在COCO和其他基准数据集上优于现有的实时检测器，例如YOLO及其变体。代码和模型可在https://github.com/Atten4Vis/LW-DETR获取。

1 介绍

实时目标检测是视觉识别中的一个重要问题，并具有广泛的实际应用。目前主流的解决方案基于卷积网络，例如YOLO系列。最近，基于Transformer的方法（(DETR)取得了显著进展。不幸的是，DETR用于实时检测尚未完全开发，其性能是否可以与最先进的卷积方法相比仍不清楚。在本文中，我们构建了一种适用于实时目标检测的轻量级DETR方法。该架构非常简单：一个普通的ViT编码器和一个通过卷积投影器连接的DETR解码器。我们提出聚合编码器中的多层特征图、中间和最终特征图，形成更强的编码特征图。我们的方法利用了有效的训练技术。例如，我们使用变形交叉注意力形成解码器、IoU感知分类损失及编码器-解码器预训练策略。

另一方面，我们的方法也利用了推理高效的技巧。例如，我们采用交错窗口和全局注意力机制，在普通的ViT编码器中用窗口注意力替换部分全局注意力以降低复杂度。我们通过窗口主特征图组织方法对交错注意力进行高效实现，有效地减少了昂贵的内存置换操作。

图1显示，所提出的简单基线在COCO上意外地超越了之前的实时检测器，例如YOLO-NAS、YOLOv8和RTMDet。这些模型通过在Objects3655上的预训练得到了改进，并且端到端的时间成本（包括NMS时间）是在官方实现设置下测量的。

我们进行了广泛的实验来与现有的实时检测算法进行比较。我们进一步优化了NMS设置，并提高了现有算法的性能。尽管如此，所提出的基线仍然优于这些算法（图1中标记为“*”）。此外，我们在更多的检测基准测试中展示了所提出方法的实验结果。

所提出的基线仅仅探索了简单且易于实现的技术，并显示出有希望的性能。我们认为我们的方法有可能从其他设计中受益，例如高效的多尺度特征融合、token稀疏化、蒸馏以及其他训练技术，如YOLO-NAS中使用的技术。我们还表明，所提出的方法适用于带有卷积编码器的DETR方法，例如ResNet-18和ResNet-50，并且实现了良好的性能。

2 相关工作

实时物体检测。实时物体检测具有广泛的实际应用。现有的最先进的实时检测器，如YOLO-NAS、YOLOv8和 RTMDet相比YOLO的第一版已经有了显著改进，这得益于检测框架、架构设计、数据增强、训练技术以及损失函数的进步。这些检测器基于卷积。在本文中，我们研究了用于实时检测的基于Transformer的解决方案，这一领域尚未得到充分探索。

ViT用于物体检测。视觉Transformer (ViT)在图像分类中表现出色。将ViT应用于物体检测通常利用窗口注意力或分层架构来减少内存和计算成本。UViT采用渐进式窗口注意力机制。ViTDet]实现了预训练的纯ViT，并结合了交错的窗口和全局注意力。我们的方法遵循ViTDet使用交错的窗口和全局注意力，并且额外采用了按窗口优先顺序组织特征图来减少内存重排的成本。

DETR及其变体。检测Transformer (DETR)是一种端到端的检测方法，它去除了许多手工设计组件的必要性，例如锚点生成和非极大值抑制(NMS)。有许多后续方法致力于改进DETR，包括架构设计、目标查询设计、训练技术以及损失函数改进。此外，还有多种工作致力于通过架构设计、计算优化、剪枝和蒸馏来降低计算复杂度。本文的兴趣在于构建一个简单的DETR基线用于实时检测，而这些方法尚未探索这一领域。

与我们的工作同步进行的是RT-DETR，该方法也将DETR框架应用于构建实时检测器，重点放在形成编码器的CNN主干上。目前的研究主要集中在相对较大的模型上，缺乏对小型模型的研究。我们的LW-DETR探索了使用纯ViT主干和DETR框架进行实时检测的可行性。

3 LW-DETR

3.1 架构

LW-DETR由一个ViT编码器、一个投影器和一个DETR解码器组成。

编码器。我们采用ViT作为检测任务的编码器。一个基础的ViT包含了一个补丁化层和多个Transformer编码层。初始版本中的一个Transformer编码层包含一个覆盖所有标记的全局自注意力层和一个前馈神经网络(FFN)层。全局自注意力层的计算成本较高，其时间复杂度与标记数量呈二次关系。我们在一些Transformer编码层中使用了窗口自注意力机制来降低计算复杂度（详见第3.4节）。我们提出聚合多级特征图，即编码器中的中间特征图和最终特征图，以形成更强的编码特征图。图2展示了编码器的一个示例。

解码器。解码器是由多个Transformer解码层堆叠而成。每一层包含自注意力层、交叉注意力层和FFN层。为了提高计算效率，我们采用了可变形交叉注意力机制。DETR及其变体通常使用6个解码层。在我们的实现中，使用了3个Transformer解码层。这使得时间消耗从1.4毫秒减少到了0.7毫秒，相对于我们的方法中tiny版本其余部分的时间成本1.3毫秒来说，这是一个显著的降低。

我们采用了一种混合查询选择方案来形成目标查询，这是内容查询和空间查询的组合。内容查询是可学习的嵌入向量，类似于DETR中的做法。空间查询则基于两阶段方案：从投影器的最后一层中选取前K个特征，预测边界框，并将对应的边界框转换为嵌入向量作为空间查询。

投影器。我们使用一个投影器来连接编码器和解码器。投影器的输入来自编码器聚合后的编码特征图。投影器是一个C2f块（跨阶段部分密集连接网络的一种扩展），在YOLOv8中有实现。

在构建LW-DETR的大尺寸和超大尺寸版本时，我们修改了投影器以输出两个尺度（1/8 和 1/32）的特征图，并相应地使用了多尺度解码器。投影器包含两个并行的C2f块。其中一个处理通过反卷积上采样获得的1/8特征图，另一个处理通过步幅卷积下采样获得的1/32特征图。图3展示了单尺度投影器和多尺度投影器的流程。

目标函数。我们使用了IoU-aware分类损失，IA-BCE损失。

其中和分别是正样本和负样本的数量。s是预测的分类得分。t是结合了IoU得分u（与真实值相比）的目标得分：，并且经验性地设置为0.25。

总体损失是分类损失和边界框损失的组合，这部分与DETR框架中的相同，其公式如下：

其中和分别设置为 2.0 和 5.0。和分别是指用于框回归的广义IoU（GIoU）损失和 L1 损失。

3.2 实例化

我们实例化了五个实时检测器：tiny、small、medium、large 和 xlarge。具体设置见表1。

tiny检测器包含一个具有6层的Transformer编码器。每一层包含一个多头自注意力模块和一个前馈神经网络（FFN）。每个图像块线性映射到一个192维的表示向量。投影器输出256通道的单尺度特征图。解码器有100个物体查询。

small检测器包含10个编码器层，并且有300个物体查询。与tiny检测器相同的是，输入块表示和投影器输出的维度均为192和256。medium检测器与small相似，不同之处在于输入块表示的维度为384，相应地，编码器的维度也为384。

large检测器包含一个10层的编码器，并使用双尺度特征图（参见第3.1节中的投影器部分）。输入块表示和投影器输出的维度均为384。xlarge检测器与large类似，不同之处在于输入块表示的维度为768。

3.3 有效的训练

更多的监督。已经开发出多种技术来引入更多监督以加速DETR的训练。我们采用了易于实现且不改变推理过程的Group DETR。遵循该方法，我们在训练过程中使用了13个并行权重共享的解码器。对于每个解码器，我们从投影器的输出特征中为每个组生成物体查询。我们在推理时使用主解码器。

在Objects365上的预训练。预训练过程分为两个阶段。首先，我们使用MIM方法CAEv2对ViT在Objects365数据集上进行预训练，该方法基于预训练模型，这在COCO数据集上带来了0.7 mAP的提升。其次，我们遵循的方法重新训练编码器并在监督方式下训练投影器和解码器。

3.4 高效推理

我们进行了简单的修改，采用交错窗口注意力和全局注意力：即用窗口自注意力层替换部分全局自注意力层。例如，在一个6层的ViT中，第一、第三和第五层使用窗口注意力实现。窗口注意力通过将特征图分割成不重叠的窗口，并分别对每个窗口执行自注意力操作来实现。

我们采用窗口优先的特征图组织方案以实现高效的交错注意力，该方案按窗口组织特征图。ViTDet的实现按照逐行（行优先组织）组织特征图，需要昂贵的转置操作以将特征图从行优先组织转换为窗口优先组织来进行窗口注意力。我们的实现在推理过程中移除了这些操作，从而减少了模型延迟。

我们通过一个示例来说明窗口优先的方式。假设有一个4 × 4的特征图：

对于窗口大小为 2 × 2 的情况，其窗口优先组织方式如下：

这种组织方式适用于窗口注意力机制和全局注意力机制，并且不需要重新排列特征。按照行优先组织方式：

这种方式适用于全局注意力机制，并且在执行窗口注意力机制时需要进行代价高昂的转置操作。

3.5 实证研究

我们通过实证展示了有效的训练和高效的推理技术如何改进 DETR。这里以小型检测器为例进行说明。本研究基于一个初始检测器：编码器在所有层中使用全局注意力机制，并输出最后一层的特征图。结果如表2所示。

延迟改进

ViTDet 采用的交错窗口注意力和全局注意力机制将计算复杂度从 23.0 GFlops 减少到 16.6 GFlops，验证了替换成本较高的全局注意力机制为成本较低的窗口注意力机制带来的好处。然而，延迟并未减少，甚至增加了 0.2 毫秒。这是因为在行优先的特征图组织方式下需要额外的代价高昂的转置操作。而窗口优先的特征图组织方式减轻了这一副作用，并导致延迟减少了 0.8 毫秒，从 3.7 毫秒减少到 2.9 毫秒。

性能改进

多级特征聚合带来了 0.7 mAP 的提升。IoU感知分类损失和更多的监督将 mAP 分数从 34.7 提高到了 35.4 和 38.4。用于框回归目标的边界框重新参数化（详见补充材料）带来了轻微的性能提升。显著的改进来自于在 Objects365 上的预训练，并达到了 8.7 mAP 的提升，这表明Transformer确实受益于大数据量。更长的训练计划可以带来进一步的改进，从而形成了我们的 LW-DETR-small 模型。

4 实验

4.1 设置

数据集 预训练的数据集为 Objects365。我们将训练集中的图像与验证集中的图像（除了前 5k 张图像）结合起来用于检测预训练。我们使用标准的 COCO2017数据划分策略，并在 COCO val2017 上进行评估。

数据增强 我们采用了 DETR 及其变体中的数据增强方法。我们遵循实时检测算法，随机将图像调整为正方形进行训练。为了评估性能和推理时间，我们按照实时检测算法中使用的评估方案，将图像调整为 640 × 640 的大小。我们使用 10 × 10 的窗口大小来确保图像尺寸可以被窗口大小整除。

实现细节 我们在 Objects365上对检测模型进行了 30 个周期的预训练，并在 COCO上进行了总共 180K 次迭代的微调。我们采用了指数移动平均 (EMA) 技术，衰减率为 0.9997。我们使用 AdamW 优化器进行训练。对于预训练，我们将投影器和 DETR 解码器的初始学习率设置为，ViT 主干网络的初始学习率设置为，批量大小为 128。对于微调，我们将投影器和 DETR 解码器的初始学习率设置为，ViT 主干网络的初始学习率设置为。我们在 tiny、small 和 medium 模型中设置批量大小为 32，在 large 和 xlarge 模型中设置批量大小为 16。tiny、small 和 medium 模型的训练迭代次数为 180K，相当于 50 个周期；而对于 large 和 xlarge 模型，则为 25 个周期。更多细节，例如权重衰减、ViT 编码器中的逐层衰减以及微调过程中的组件级衰减等，均在补充材料中有详细说明。

我们在 COCO val2017 上使用 T4 GPU，以 fp16 精度和批量大小为 1 的方式测量端到端的平均推理延迟，环境配置包括 TensorRT-8.6.1、CUDA-11.6 和 CuDNN-8.7.0。TensorRT 中的 efficientNMSPlugin 用来进行非极大值抑制 (NMS) 的实现。我们使用官方实现对所有实时检测器的性能和端到端延迟进行了测量。

4.2 结果

我们的五个 LW-DETR 模型的结果如表 3 所示。LW-DETR-tiny 在 T4 GPU 上实现了 42.6 mAP 和 500 FPS。LW-DETR-small 和 LW-DETR-medium 分别达到了 48.0 mAP 和超过 340 FPS，以及 52.5 mAP 和超过 178 FPS 的速度。大型和超大型模型分别实现了 56.1 mAP 和 113 FPS，以及 58.3 mAP 和 52 FPS。

与先进实时检测器的比较 在表 3 中，我们报告了 LW-DETR 模型与代表性实时检测器（包括 YOLO-NAS、YOLOv8和 RTMDet）的比较结果。可以看出，无论是否使用预训练，LW-DETR 均一致优于先前的最佳实时检测器。在 tiny 到 xlarge 的五种尺度下，我们的 LW-DETR 在延迟和检测性能方面明显优于 YOLOv8 和 RTMDet。

与之前最佳方法之一 YOLO-NAS 相比，该方法通过神经架构搜索获得，我们的 LW-DETR 在 small 和 medium 尺度上分别高出 0.4 mAP 和 0.9 mAP，并且运行速度分别快 1.6 倍和约 1.4 倍。当模型增大时，改进更加显著：在 large 尺度下以相同速度运行时，提高了 3.8 mAP。

我们进一步通过在 NMS 过程中精细调整分类得分阈值来改进其他方法，并在右侧两列报告了结果。结果得到了显著提升，但仍低于我们的 LW-DETR。我们预计，我们的方法还可以从其他改进中受益，例如神经架构搜索 (NAS)、数据增强、伪标记数据和知识蒸馏等，这些方法已被先前的实时检测器使用。

与同期工作的比较 我们还将在实时检测领域的同期工作（YOLO-MS、Gold-YOLO、RT-DETR 和 YOLOv10）与我们的 LW-DETR 进行比较。YOLO-MS 通过增强多尺度特征表示来提高性能。Gold-YOLO 提升了多尺度特征融合，并应用了 MAE 风格的预训练来改善 YOLO 性能。YOLOv10 设计了几种效率和准确性驱动的模块来提高性能。RT-DETR与 LW-DETR 密切相关，基于 DETR 框架构建，但在主干、投影器、解码器和训练方案方面与我们的方法存在许多差异。

比较结果如表 4 和图 4 所示。我们的 LW-DETR 在检测性能和延迟之间始终取得了更好的平衡。YOLO-MS 和 Gold-YOLO 在所有模型尺度上均显示出比 LW-DETR 更差的结果。LW-DETR-large 比密切相关的 RT-DETR-R50 提高了 0.8 mAP 并显示更快的速度（8.8 ms vs. 9.9 ms）。其他尺度下的 LW-DETR 也表现出优于 RT-DETR 的结果。与最新工作 YOLOv10-X相比，我们的 LW-DETR-large 达到了更高的性能（56.1 mAP vs. 54.4 mAP），并且延迟更低（8.8 ms vs. 10.70 ms）。

4.3 讨论

NMS 后处理 DETR 方法是一种端到端算法，不需要非极大值抑制（NMS）后处理步骤。相比之下，现有的实时检测器（如 YOLO-NAS、YOLOv8和 RTMDet）需要 NMS后处理。NMS 程序会额外消耗时间。我们将这部分额外时间计算在内，以测量端到端推理成本，这对于实际应用非常重要。图 1 和表 3 展示了官方实现中使用 NMS 设置的结果。

我们进一步通过调整 NMS 后处理中的分类得分阈值来改进带有 NMS 的方法。我们发现默认的得分阈值（0.001）在 YOLO-NAS、YOLOv8 和 RTMDet 中导致较高的 mAP，但产生了大量框体，从而增加了延迟。特别是对于小型模型，端到端延迟主要由 NMS 延迟主导。通过调整阈值，我们获得了 mAP 得分与延迟之间的良好平衡。观察结果显示，mAP 分数略有下降（例如 -0.1 至 -0.5 mAP），但运行时间大幅减少（例如 RTMDet 和 YOLOv8 减少了 4 至 5 毫秒，YOLO-NAS 减少 1 至 2 毫秒）。这些减少是因为调整得分阈值后，更少的预测框进入 NMS。不同得分阈值的详细结果及其在 COCO val2017 数据集上的框体数量分布情况在补充材料中有给出。图 1 显示了与其他带有优化 NMS 程序的方法的对比。带有 NMS 的方法得到了改进，而我们的方法仍然优于其他方法。排名第二的方法 YOLO-NAS 是一种网络架构搜索算法，其表现接近于所提出的基线。我们认为复杂的网络架构搜索过程可能对 DETR 方法有益，并且可以期待进一步改进。

预训练 我们经验性地研究了预训练的效果。表 5 中的结果表明，预训练对我们的方法有显著改进，平均提升了 5.5 mAP。其中 tiny 模型获得了 6.1 mAP 的提升，xlarge 模型则获得了 5.3 mAP 的提升。这表明在大规模数据集上的预训练对基于 DETR 的模型非常有益。我们进一步展示了训练过程适用于带有卷积编码器的 DETR 方法。我们用 ResNet-18 和 ResNet-50 替换了变压器编码器。如表 5 所示，这些 LW-DETR 变体在延迟和 mAP 方面与带有变压器编码器的 LW-DETR 接近，而且预训练带来的益处与后者相似但略低一些。

同时，我们也研究了非端到端检测器的预训练改进。根据表 3、表 4 和表 6 的结果，似乎仅在 Objects365 上的预训练对于非端到端检测器的收益有限，这与基于 DETR 的检测器中的现象不同，在后者中预训练带来了巨大改进。鉴于非端到端检测器（如 YOLOv8）训练 300 甚至 500 个周期的情况，我们怀疑有限的增益是否与训练周期有关。我们比较了随着训练周期增加，预训练权重带来的改进。表 6 显示，随着训练周期的增加，改进逐渐减少，这在一定程度上支持了上述假设。上述说明仅是初步步骤，我们认为还需要更多的研究来揭示预训练效益差异的根本原因。

4.4 在更多数据集上的实验

跨域评估 评估模型泛化能力的一种方法是在不同域的数据集上直接评估模型。我们采用一个类别无关的目标检测基准——UVO，该数据集中 57% 的目标实例不属于 COCO 的 80 个类别中的任何一个。UVO 基于 YouTube 视频，其外观与 COCO 大不相同，例如，某些视频采用第一人称视角并且存在明显的运动模糊。我们在 UVO 验证集上评估了在 COCO 上训练的实时检测器（取自表 3）。表 7 提供了评估结果。LW-DETR 在跨域评估中优于竞争的最新实时检测器。具体来说，LW-DETR-small 在 mAP 和 AR 方面分别比 RTMDet-s、YOLOv8-s 和 YOLO-NAS-s 的最佳结果高出 1.3 mAP 和 4.1 AR。在召回率方面，它也表现出更强的能力，能够在不同尺度（小、中、大）下检测更多目标。上述发现表明，LW-DETR 相较于先前实时检测器的优势不仅在于针对 COCO 的特定调优，还在于其生成更通用模型的能力。

多域微调 另一种方法是在不同域的小型数据集上对预训练检测器进行微调。RF100 包含 100 个小数据集，涵盖 7 种图像域，包含 224,000 张图像和 829 个类别标签。这有助于研究人员利用真实数据测试模型的泛化能力。我们在 RF100 的每个小型数据集上对实时检测器进行了微调。表 8 给出了微调结果。LW-DETR-small 在不同域上均优于当前最新的实时检测器。特别是在“文档”和“电磁”领域，我们的 LW-DETR 显著优于 YOLOv5、YOLOv7、RTMDet 和 YOLOv8（分别高出 5.7 AP 和 5.6 AP）。LW-DETR-medium 整体上还能带来进一步的改进。这些发现突显了 LW-DETR 的多功能性，使其成为一系列封闭域任务的强大基线。