【魔改YOLO系列】Mamba YOLO: 基于SSMs的YOLO用于目标检测

科技   2024-11-06 10:06   辽宁  

点击下方深度学习爱好者”,选择加"星标"或“置顶

论文信息

题目:Mamba YOLO: SSMs-Based YOLO For Object Detection

Mamba YOLO: 基于SSMs的YOLO用于目标检测

作者:Zeyu Wang, Chen Li, Huiying Xu, Xinzhong Zhu

源码:https://github.com/HZAI-ZJNU/Mamba-YOLO

论文创新点

  1. 基于SSM的新型目标检测模型:作者提出了Mamba-YOLO,这是一种基于状态空间模型(SSM)的新型目标检测模型。该模型不仅优化了SSM基础,还特别适应于目标检测任务,为YOLO系列在目标检测中建立了新的基准。
  2. ODSSBlock模块的创新设计:为了弥补SSM在局部建模能力的不足,作者设计了ODSSBlock模块,其中包括LS Block有效地提取输入特征图的局部空间信息,以及RG Block通过结合门控聚合与有效卷积和残差连接,增强模型对局部依赖性的捕获能力。
  3. 不同规模模型的设计:作者设计了一套不同规模的模型**Mamba-YOLO (Tiny/Base/Large)**,以支持不同大小和规模的任务部署,这显示了模型的灵活性和广泛的适用性。

摘要

随着深度学习技术的快速发展,YOLO系列为实时目标检测器树立了新的基准。研究者们在YOLO的基础上不断探索重参数化、高效层聚合网络和无锚点技术的新应用。为了进一步提升检测性能,引入了基于Transformer的结构,显著扩大了模型的感受野并取得了显著的性能提升。然而,这种改进的代价是自注意力机制的二次复杂度增加了模型的计算负担。幸运的是,状态空间模型(SSM)作为一种创新技术的出现有效地缓解了由二次复杂度引起的问题。鉴于这些进展,我们引入了Mamba-YOLO,一种基于SSM的新型目标检测模型。Mamba-YOLO不仅优化了SSM基础,还特别适应于目标检测任务。鉴于SSM在序列建模中的潜在局限性,如感受野不足和图像局部性弱,我们设计了LSBlock和RSBlock。这些模块能够更精确地捕获局部图像依赖性,并显著增强模型的鲁棒性。在公开可用的基准数据集COCO和VOC上的广泛实验结果表明,Mamba-YOLO在性能和竞争力方面超过了现有的YOLO系列模型,展示了其巨大的潜力和竞争优势。

关键词

YOLO系列,目标检测,状态空间模型(SSM),Mamba-YOLO,实时检测

3 方法

总体架构

Mamba YOLO的架构概述如图2所示。我们的目标检测模型分为ODMamba骨干和颈部部分。ODMamba由Simple Stem和Downsample Block组成。在颈部,我们遵循PAN-FPN的设计,使用ODSSBlock模块代替C2f[42]以捕获更丰富的梯度信息流。骨干首先通过Stem模块进行下采样,得到分辨率为的2D特征图。随后,所有模型由ODSSBlock组成,后面跟着VisionClue Merge模块以进一步下采样。在颈部部分,我们采用PAFPN[42]的设计,使用ODSSBlock替换C2f,其中Conv仅负责下采样。
Simple Stem现代视觉Transformer(ViTs)通常采用分割的补丁作为初始模块,将图像划分为不重叠的段。这种分割过程是通过卷积操作实现的,核大小为4,步长为4。然而,最近的研究,如EfficientFormerV2[43],表明这种方法可能会限制ViTs的优化能力,影响整体性能。为了在性能和效率之间取得平衡,我们提出了一个简化的茎层。与使用不重叠的补丁不同,我们采用两个卷积,步长为2,核大小为3。
Vision Clue Merge虽然卷积神经网络(CNN)和视觉Transformer(ViT)结构通常采用卷积进行下采样,但我们发现这种方法会干扰SS2D[31]在不同信息流阶段的选择性操作。为了解决这个问题,VMamba[31]将2D特征图拆分,并使用卷积降低维度。我们的发现表明,为状态空间模型(SSMs)保留更多的视觉线索有利于模型训练。与传统的维度减半方法不同,我们通过以下方式简化这个过程:1)去除规范;2)拆分维度图;3)将多余的特征图附加到通道维度;4)使用3 \times 3$卷积不同,我们的方法保留了SS2D从前一层选择的特征图。

ODSS Block

如图4所示,ODSS Block是Mamba YOLO的核心模块,在输入阶段经历了一系列处理,使网络能够学习更深入、更丰富的特征表示,同时通过批量归一化保持训练推理过程的高效和稳定。
其中表示激活函数(非线性SiLU)。ODSS Block的层归一化和残差链接设计借鉴了Transformer Blocks[6]风格的架构,允许模型在深层堆叠的情况下高效流动和训练。计算公式如下:
其中分别表示输入和输出特征,LS和RG分别表示LocalSpatial Block和ResGated Block,表示经过2D-Selective-Scan[31]后的中间状态。Scan Expansion、S6 Block和Scan Merge是SS2D[31]算法的三个主要步骤,其主要流程如图3所示。扫描扩展操作将输入图像扩展为一系列子图像,每个子图像表示一个特定方向,从对角线视角观察,扫描扩展操作沿四个对称方向进行,分别是自上而下、自下而上、左至右和右至左。这种布局不仅全面覆盖了输入图像的所有区域,还通过提供丰富的多维信息基础,增强了多维捕获图像特征的效率和全面性,为后续通过系统的方向变换进行特征提取提供了基础。然后这些子图像在S6 block[32]操作中进行特征提取,最后通过扫描合并操作,这些子图像合并在一起,形成与输入图像大小相同的输出图像。

LocalSpatial Block

Mamba架构已被证明在捕获长距离地面依赖性方面是有效的。然而,在处理涉及复杂尺度变化的任务时,它在提取局部特征方面面临一定的挑战。在图4(c)中,本文提出了LocalSpatial Block以增强局部特征的捕获。具体来说,对于给定的输入特征,它首先经历深度可分离卷积,该卷积独立操作每个输入通道,而不混合通道信息。有效地提取输入特征图的局部空间信息,同时减少计算成本和参数数量,然后经过批量归一化以提供一定程度的正则化效果,减少过拟合,得到的中间状态定义为:
中间状态通过卷积混合通道信息,并通过对激活函数的使用更好地保留信息分布,使模型能够学习更复杂的特征表示,能够从输入特征图中提取丰富的多尺度上下文信息。在LSBlock中,激活函数使用非线性GELU来改变特征的通道数,而不改变空间维度,从而增强特征表示。最后,通过残差连接将原始输入与处理后的特征融合。使模型能够理解和整合图像中不同维度的特征,从而增强对尺度变化的鲁棒性。
其中是输出特征,表示激活函数。LSBlock能够有效地捕获和表示输入特征图的局部空间信息,并将其与原始输入融合,增强特征的表示,为Mamba YOLO提供了强大的能力来处理视觉数据中的尺度变化和上下文信息。

ResGated Block

原始的MLP仍然是最广泛采用的,VMamba[31]架构中的MLP也遵循Transformer设计,对输入序列执行非线性变换,以增强模型的表达能力。近期的研究,Gated MLP[44, 45]在自然语言处理中显示出强大的性能,我们发现门控机制在视觉领域同样具有潜力。在图4(d)中,本文提出ResGated Block旨在以低成本提高模型性能,RG Block从输入创建两个分支,并在每个分支上以卷积的形式实现全连接层。
的分支上使用深度可分离卷积作为位置编码模块,通过残差连接的方式在训练期间更有效地回传梯度,这降低了计算成本,并显著提高了性能,同时保留了和利用了图像的空间结构信息。RG Block采用非线性GeLU作为激活函数来控制每个级别的信息流,然后通过元素乘法与的一个分支合并,然后通过卷积与全局特征混合以融合通道信息,最后通过残差连接与原始输入在隐藏层的特征相加。RG Block在只带来轻微的计算成本增加的同时,能够捕获更多的全局特征,产生的输出特征定义为:
其中表示激活函数(非线性GeLU)。在本文中,RG Block中的门控机制通过集成卷积操作保留空间信息,使模型对图像中的细粒度特征更加敏感。与传统的MLP相比,RG Block将全局依赖性和全局特征传递到每个像素,以捕获邻近特征的依赖性,这使得上下文信息丰富,进一步增强了模型的表达能力。

4 实验

声明

本文内容为论文学习收获分享,受限于知识能力,本文对原文的理解可能存在偏差,最终内容以原论文为准。本文信息旨在传播和学术交流,其内容由作者负责,不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题,请及时与作者联系,作者将在第一时间回复并处理。
下载1:Pytoch常用函数手册
在「深度学习爱好者」公众号后台回复:Pytorch常用函数手册,即可下载全网第一份Pytorch常用函数手册,涵盖Tensors介绍、基础函数介绍、数据处理函数、优化函数、CUDA编程、多线程处理等十四章章内容。
下载2:Python视觉实战项目52讲
小白学视觉公众号后台回复:Python视觉实战项目即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目,助力快速学校计算机视觉。

交流群

欢迎加入公众号读者群一起和同行交流,目前有SLAM、三维视觉驶、计算摄影、检测、分割、识别、医学影像、GAN算法竞赛等微信群(以后会逐渐细分),请扫描下面微信号加群,备注:”昵称+学校/公司+研究方向“,例如:”张三 + 上海交大 + 视觉SLAM“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告,否则会请出群,谢谢理解~

深度学习爱好者
分享机器学习、深度学习和Python等知识与原理,每天分享深度学习与计算机视觉领域的经典和最新的论文总结,带读者一起跟踪前言科技!
 最新文章