目标检测、实例分割、旋转框样样精通!详解高性能检测算法 RTMDet

语言: CN / TW / HK

1. 简介

近几年来,目标检测模型,尤其是单阶段目标检测模型在工业场景中已经得到广泛应用。对于检测算法来说,模型的精度以及运行效率是实际使用时最受关注的指标。因此,我们对目前的单阶段目标检测器进行了全面的改进:从增强模型的特征提取能力和对各个组件的计算量进行均衡化,到使用动态的软标签来优化训练策略,再到对数据增强的效率和性能进行改进, 我们对算法的各个方面都提出了新的优化方案。

通过这些改进,我们得到了从 tiny 到 extra-large 大小的一系列高性能检测模型,我们将这套实时目标检测模型(R eal- T ime M odels for object Det ection)命名为 RTMDet。其中,RTMDet-x 取得了 52.8 mAP 的精度和 300+FPS 的运行速度,同时,RTMDet-tiny 在仅有 4M 参数量的情况下达到了 41.1 mAP 的精度,超越了同级别的所有模型。RTMDet 不仅仅在目标检测这一任务上性能优异,在实时实例分割以及旋转目标检测这两个任务中也同样达到了 SOTA 的水平!

除了拥有超高的精度,RTMDet 的工程优化也非常全面:基于 MMDeploy,RTMDet 全系列已经支持了 ONNXRuntime 以及 TensorRT 等推理框架的部署。同时,为了节约广大炼丹师的宝贵时间,RTMDet 的训练效率也同样进行了大幅的优化,我们在 MMYOLO 中提供了高效训练版本的代码,RTMDet-s 训练 300 epoch 仅需 14 × 8 GPU 时!

检测和实例分割模型代码已开源至:

https://github.com/open-mmlab/mmdetection/tree/3.x/configs/rtmdet

基于 MMYOLO 的高效训练版本:

https://github.com/open-mmlab/mmyolo/tree/dev/configs/rtmdet

基于 MMRotate 的旋转框检测模型:

https://github.com/open-mmlab/mmrotate/tree/1.x/configs/rotated_rtmdet

技术报告链接:

https://arxiv.org/abs/2212.07784

话不多说,接下来就让我们来详细了解 RTMDet 在各方面的改进吧~

2. 算法和工程优化详解

2.1 模型结构设计

拥有一个高效的模型结构是设计实时目标检测器最关键的问题之一。自 YOLOv4 将 Cross Stage Partial Network 的结构引入 DarkNet 之后,CSPDarkNet 因其简洁高效而被广泛应用于 YOLO 系列的各个改进版中。

而 RTMDet 也将 CSPDarkNet 作为基线,并使用同样构建单元组成的 CSPPAFPN 进行多尺度的特征融合,最后将特征输入给不同的检测头,进行目标检测、实例分割和旋转框检测等任务。整体的模型结构如下图所示:

图 1. RTMDet 模型结构图

整体的宏观架构与 YOLO 系列并没有太大区别,RTMDet 对模型的修改主要聚焦于两点:增加基本构建单元的特征提取能力,以及探索模型不同组件间的计算效率和精度的平衡。

首先让我们回顾一下 DarkNet 基本的构建单元,如图 2.(a) 所示,它由 1 个 1x1 的卷积和 1 个 3x3 的卷积组成,以及残差连接组成。为了提升基本单元的特征提取能力,我们选择引入更大的卷积核来增大感受野。在经过对不同 kernel size 的消融实验(表 1)后我们发现,在基本单元中加入一个 5x5 的深度可分离卷积,取得了最好的计算效率与精度的平衡。新的基础单元结构如 图 2.(b) 所示。

表 1. 不同 kernel size 的精度与速度的对比实验

值得注意的是,最近提出的 YOLO 的不同改进版,如 YOLOv6,v7,以及 PPYOLO-E,都选择在基本单元中加入 RepVGG 的重参数化卷积模块,如图 2.(c),图 2.(d) 所示。

图 2. CSPDarknet、RTMDet 、PPYOLOE、YOLOv6 的基本构建单元

重参数化通过训练时使用多分支的结构,然后再在推理时融合成单分支的方式,在不增加推理计算量的情况下提升性能。

然而,重参数化也有其弊端,比如会明显增加训练的显存开销,同时也会导致使用低精度对模型进行量化时产生较大的性能下降。尽管可以通过 QAT 等方式解决量化掉点的问题,但这也使得从模型训练到部署的流程变得更为复杂。

因此,我们认为引入大卷积核增加感受野的方式相比引入重参数化来说,不论是从训练还是部署都能够提供更好的效果。

除了基本构建单元的修改之外,RTMDet 还对整个模型的不同分辨率层级之间、以及 backbone 和 neck 之间的计算量分配上进行了全面的调整。由于在基本单元中引入了深度可分离卷积,使得模型整体的层数相比于 CSPDarkNet 来说变得更深,从而导致推理速度的减慢。

为了解决这一问题,我们对模型不同分辨率层级的基本单元数量进行了调整,从原本 C2~C5 分别为 3-9-9-3 个 block,调整为了 3-6-6-3 个 block

图 3. Backbone 结构改进

同时为了保持模型的整体计算量不变,我们略微增大了模型的宽度,并且在每个 stage 之后增加一个 ChannelAttention 模块来提供通道注意力(如图 4)。

图 4. 增加通道注意力前后的 C5 stage 特征对比

如表 2 所示,在经过调整之后,模型的推理速度更快,并且能够保持相似的精度

表 2. Backbone 结构对比实验

而对于用来进行多尺度特征融合的 neck 模块来说,以往的方法,如 EfficientDet、GiraffeDet 等,都是通过增加更多的特征融合次数,也就是加入更多的连接来提升性能。

然而,过多的特征层之间的连接会显著增加显存的占用,也会因为访存问题而增加推理耗时。因此,我们选择仅通过增加 neck 模块的计算量占比来提升其性能,当 backbone 与 neck 的参数量调整至相似时,整个检测器的推理速度更快,性能也更高:

表 3. Backbone neck 参数量分配对比实验

除了 backbone 和 neck 之外,我们对模型的检测头也进行了调整。YOLO 系列往往在不同的特征层级上分别使用独立的检测头,这就导致了模型参数的利用率低下。因为对于不同的特征层级来说,其检测出的物体的特征,在相对的尺度大小下应当是相近的,而学术界常用的检测器(如 RetinaNet、FCOS)使用共享参数的检测头也正印证了这一点。

但是,由于不同层级之间特征的统计量仍存在差异,Normalization layer 依然是必须的,由于直接在共享参数的检测头中引入 BN 会导致其滑动平均值产生误差,而引入 GN 又会增加推理时的开销,因此我们参考 NASFPN 的做法,让检测头共享卷积层,而 BN 则分别独立计算

在使用共享的检测头后,模型的参数量得到了减少,而且性能非但没有下降,反而还得到了略微的提升:

表 4. Head 结构对比

通过采用以上这些策略,我们又对通道和深度的超参数进行缩放,得到了 tiny、s、m、l、x 五种不同大小的模型,其参数量分别从 4M 到 90M 递增,以提供给不同的应用场景。

2.2 训练策略优化

正负样本的标签分配策略是目标检测训练过程中最重要的一环,近几年来,标签分配策略从最初的基于 anchor IoU 的静态匹配方式,逐渐演进为使用代价函数进行动态标签分配。然而,目前主流的动态标签分配策略,如匈牙利匹配、OTA 等均使用与损失函数一致的函数计算代价矩阵。

我们经过实验发现,与损失函数完全一致的代价矩阵计算方式并不是最优的。因此,RTMDet 基于 YOLOX 的 SimOTA 进行改进,使用了动态的软标签分配策略,其代价矩阵计算公式如下:

它由三个代价函数组成,首先是分类代价函数,传统的分类代价往往使用 0-1 分布的二值化标签进行损失计算,这很容易导致,一个拥有很高分类置信度但是拥有错误检测框的低质量预测结果得到一个很低的分类代价,反之也是如此。

因此我们参考 GFL,将预测框与 Ground Truth 的 IoU 得分作为软标签,并对不同得分的匹配进行了重新加权,使得分类代价的匹配结果更为准确和稳定,其公式如下:

而对于回归代价,使用与损失函数一致的 GIoU,很容易导致一个低质量的匹配和一个高质量的匹配区分度不够高,这是因为对于 GIoU 函数来说,一个完全正确的检测框和一个完全错误的检测框,他们之间的 IoU 得分也只相差了 1 而已。因此,我们选择取对数来增大低质量与高质量匹配之间的差异

最后,我们还引入了一个“软化”的中心先验代价。与之前的众多方法使用固定的先验区域(如 ATSS,SimOTA)不同,中心先验代价将位置先验也引入代价矩阵的计算之中,一方面能够稳定收敛,另一方面也能够使匹配的正样本区域更为灵活。

为了公平比较,我们在标准的 ResNet-50 12 epoch 的 setting 下与其他标签分配策略进行了对比,结果显示我们的方法取得了最优的精度:

表 5. R50 1x setting 下的标签分配策略性能对比

为了验证这套方法的通用性,我们也在 300 epoch 和强数据增强的情况下与 YOLOX 的 SimOTA 进行了对比,结果显示,我们的方法也同样得到了更高的精度:

表 6. 使用 RTMDet-s 训练 300 epoch 的标签分配策略性能对比

我们使用 MMYOLO 中提供的 GradCAM++ 可视化工具(demo/boxam_vis_demo.py)可视化了 neck 部分的特征图,从下图中可以看出,与 YOLOv6 相比,使用了动态软标签分配策略的 RTMDet 的检测目标的特征响应非常集中,且没有误检:

图 5. 使用 GradCAM++ 对 neck 部分特征响应的可视化(左:RTMDet-l,右:YOLOv6-l)

2.3 数据增强改进

数据增强不仅关乎模型的精度,也对训练的效率有着巨大的影响。随着 GPU 计算性能的不断增加,模型前向和反向传播的速度在不断提升。

然而,当我们在训练工业级的模型时,由于引入大量的数据增强,CPU 的计算能力以及存储的 IO 往往成为了制约模型训练速度的瓶颈。尤其是当使用了 YOLO 系列中广泛使用的 Mosaic 和 MixUp 这样涉及到多张图片混合的数据增强,由于需要读取和处理的图片数量成倍增加,数据增强的耗时也大幅增加。

图 6. Mosaic MixUp 数据增强效果图

为了解决多图混合数据增强的耗时问题,我们在这些数据增强中引入了缓存机制。我们使用一个缓存队列将历史的图片保存下来,当需要进行图片混合操作时,不再通过 dataset 重新加载图片,而是从缓存队列中随机选取历史的图片进行混合。通过实验发现,当缓存队列足够大,且使用随机出队的方式时,训练得到的模型精度与传统的 Mosaic & MixUp 并无区别。由于引入了缓存机制,两种数据增强的运行效率得到了大幅提升。

通过使用 MMDetection 中的 benchmark 工具 tools/analysis_tools/benchmark.py,我们将 RTMDet 与以高效训练著称的 YOLOv5 的数据增强进行了对比:

YOLOv5 的数据增强(MixUp 概率仅为 0.1)的吞吐量: Overall fps: 30.2 img/s, times per img: 33.1 ms/img;RTMDet 的数据增强(MixUp 全程开启)的吞吐量:Overall fps: 52.4 img/s, times per img: 19.1 ms/img。

从 benchmark 的结果来看, RTMDet 在全程开启图像混合的数据增强的情况下,数据处理吞吐量依然显著高于 YOLOv5,每秒处理的图片数量达到了 YOLOv5 的 1.7 倍!使用缓存机制加速的 Mosaic 和 MixUp 已经在 MMYOLO 中得到了全面的支持,不仅 RTMDet 可以使用,YOLOv5,v6,v7上都可以通过 use_cached=True 开关直接开启,真正做到了一键加速训练!

除了优化数据增强的运行效率,我们也对数据增强的逻辑进行了优化。我们分析了 YOLOX 的数据增强后发现了其中的问题:YOLOX 使用强弱两个阶段的数据增强,但其第一个训练阶段引入了旋转和切变等操作,导致标注框产生了比较大的误差,使其第二个训练阶段需要对回归分支额外增加一个 L1 loss 来对回归任务进行修正。这就造成了数据增强与模型的损失函数设计产生了耦合。

为了使数据增强和模型解耦,得到更通用的增强策略,我们选择在第一个训练阶段不引入导致标注误差的数据增强,而是通过增加 Mosaic MixUp 混合图片的数量来提升强度,我们将总体混合的图片数量从 YOLOX 使用的 5 张增加至 8 张。

得益于上文提到的缓存机制,增加混合图片的数量并不会导致训练速度变慢。在第二个训练阶段,我们关闭 Mosaic 和 MixUp,转而使用 Large Scale Jitter(LSJ),使模型在更符合原数据集特征分布的状态下微调。通过对比实验可以看出,我们的数据增强显著优于之前的方法:

表 7. 数据增强训练精度对比

2.4 最终效果

通过上文的种种改进,我们最终得到了 RTMDet 系列模型,我们可以将实验拆解,并逐个应用在我们的对照组模型 YOLOX 之上,来解析如何从 YOLOX 一步一步修改为 RTMDet

表 8. 从 YOLOX 到 RTMDet step-by-step 对比

首先,为了解决 CosineLr 学习率衰减过快以及 SGD 收敛不稳定的问题,我们将优化器以及学习率分别更换为了 AdamWFlatCosineLR,得到了 0.4% 的提升

然后我们使用新的基础单元构建而成的 backbone 和 neck 替换了原有的模型结构,这一改进又提升了 1.2% 的 AP,并且推理速度只减慢了 0.02ms。在使用共享权重的检测头后,参数量过大的问题得到了解决,而模型精度和推理速度都没有下降。在此基础之上,我们又增加了动态软标签分配策略以及改进后的数据增强,分别带来了 1.1% 和 1.3% 的精度提升

最后,锦上添花的一点,我们又将 backbone 在 imagenet 上进行了预训练,也使精度得到了略微的提升。不过预训练模型仅在 tiny 和 s 模型上有精度提升的效果,在更大的模型结构上则体现不出优势。

综合以上这些修改,RTMDet 相比于 YOLOX 在相似的模型大小以及推理速度上提升了 4.3% AP

通过调整深度以及宽度的系数进行模型缩放,我们最终得到了 tiny/s/m/l/x 五种不同大小的模型,在不同量级上均超越了同级别的模型

表 9. RTMDet 与其他同级别模型的性能对比(推理测速使用 3090 GPU,TensorRT 8.4.3)

  • 另外,我们还在工业界常用的 T4 GPU 上进行了测速,效果如下表所示:
  • | Model | latency(ms) | FPS | | ----------- | --------------- | ------- | | RTMDet-tiny | 2.34 | 427.35 | | RTMDet-s | 2.96 | 337.84 | | RTMDet-m | 6.41 | 156.01 | | RTMDet-l | 10.32 | 96.90 | | RTMDet-x | 18.80 | 53.19 |

表 10. RTMDet 在 T4 GPU 上的推理速度(测速使用 TensorRT 8.4,FP16,batchsize=1)

3. 多项任务取得 SOTA

为了验证 RTMDet 算法的通用性,我们通过仅增加任务头的方式,对模型进行了最小限度的修改,将其拓展至了实例分割以及旋转目标检测任务上。在仅进行 非常简单的修改的情况下,RTMDet 也依然取得了 SOTA 的效果!

图 7. RTMDet 在三种不同任务上的可视化效果

3.1 实例分割

传统的实例分割往往采用双阶段的方式来预测 mask,但近两年来,基于 kernel 的方法正逐渐兴起。为了保持单阶段检测器的简洁高效和易于部署的特性,我们参考 CondInst 为 RTMDet 增加了 mask head 以及 kernel head。

mask head 由 4 个卷积层组成,通过 neck 输出的多尺度特征预测出维度为 8 的 mask prototype feature。而 kernel head 则为每个 instance 预测 169 维的向量,从而组成 3 个动态卷积的卷积核,与 mask feature 进行交互,最终得到每个 instance 的 mask。

图 8. RTMDet-Ins 实例分割分支示意图

为了最大化利用 mask 的标注信息,我们还将标签分配中的中心先验代价的检测框中心修改为了 mask 的重心。为了与其他实例分割方法做公平对比,我们使用 ResNet50 FPN 12 epoch 的标准 setting 进行了实验,同时参考 CondInst 加入了语义分割的辅助分支加快收敛。

实验结果表明,我们的方法在标准 setting 下,不仅超越了 CondInst,SOLOv2 等单阶段模型,也超越了 Cascade Mask R-CNN 这样的多阶段模型。在使用与 RTMDet 一致的模型结构以及训练策略后,得到的 RTMDet-Ins 模型,精度不仅大幅超越 YOLOv5-seg,也超越了前几天刚刚推出的 YOLOv8,取得了实时实例分割的 SOTA

表 11. RTMDet-Ins 与其他实例分割算法的性能对比

3.2 旋转目标检测

旋转目标检测是遥感领域使用最广泛的方法,而将 RTMDet 拓展至此任务也非常简单,得益于 MMDetection 3.0 版本,我们只需要增加回归分支的输出特征维度,增加角度分量,同时更换 box 的编解码器, 便能让 RTMDet 支持预测旋转框。由于旋转框与普通目标检测任务仅有回归分支有差异,因此旋转框的模型也能够加载目标检测的预训练模型并从中受益。

我们在遥感领域最常用的 DOTA 数据集上验证了我们的方法,结果表明,我们的模型不仅超越了众多的多阶段方法也超越了之前最优的算法 PPYOLOE-R,成为了 SOTA

图 9. RTMDet-R 与其他旋转目标检测方法在 DOTA 数据集上的对比

为了验证模型的泛化能力,我们还在 DOTA 1.5 以及 HRSC 两个数据集上进行了训练,也同样取得了最优的结果

表 12. RTMDet-R 在 DOTA 1.5 上的性能

表 13. RTMDet-R 在 HRSC 上的性能,仅使用 tiny 模型就超越了之前的多阶段算法

4. 总结

通过在模型结构的基本单元引入大 kernel 深度可分离卷积增大感受野,并平衡不同分辨率层级间以及 backbone 和 neck 间的计算量、参数量平衡,改进标签分配以及数据增强策略,RTMDet 不论是在目标检测任务上,还是在实例分割以及旋转目标检测任务上,均取得了优异的性能。我们希望我们探索出的这些改进方案能够对设计实时的目标检测模型有所启发,也希望我们在工程及算法上的优化能够在实际的工业场景中有所应用。

最后,别忘了给 MMDetection、MMYOLO、MMRotate 点个 Star 哦~

目标检测和实例分割项目地址:https://github.com/open-mmlab/mmdetection/tree/3.x/configs/rtmdet

MMYOLO 高效训练版本地址:https://github.com/open-mmlab/mmyolo/tree/dev/configs/rtmdet

旋转框检测项目地址:https://github.com/open-mmlab/mmrotate/tree/1.x/configs/rotated_rtmdet