Swin-UMamba:基于ImageNet的预训练的基于Mamba的UNet网络

科技   2024-11-01 10:25   中国香港  

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

论文信息

题目:Swin-UMamba: Mamba-based UNet with ImageNet-based pretraining

Swin-UMamba:基于ImageNet的预训练的基于Mamba的UNet网络

作者:Jiarun Liu,Hao Yang,Hong-Yu Zhou,Yan Xi,Lequan Yu,Cheng Li,Yong Liang,Guangming Shi,Yizhou Yu,Shaoting Zhang,Hairong Zheng,Shanshan Wang

源码链接:https://github.com/JiarunLiu/Swin-UMamba

论文创新点

  1. 首次尝试:作者是第一次尝试发现预训练的基于Mamba的网络在医学图像分割中的影响。实验验证了基于ImageNet的预训练对于基于Mamba的网络在医学图像分割中的重要性,有时是至关重要的。
  2. 提出的模型:作者提出了两个基于Mamba的网络Swin-UMambaSwin-UMamba†,它们专门设计用于统一预训练模型的力量,并针对现实世界的部署具有不同的计算需求。这些模型展示了在各种医学图像分割数据集上优越的性能,能够以显著的优势超越以前的分割模型,包括CNN、ViT和最新的基于Mamba的模型。
  3. 基于ImageNet的预训练的影响:作者的实验表明,基于ImageNet的预训练为基于Mamba模型提供了多种优势,包括更高的分割准确性稳定的收敛性减轻过拟合问题数据效率更低的计算资源消耗。这些发现强调了预训练在提高基于Mamba模型在视觉任务中的性能和效率中的重要性。

摘要

精确的医学图像分割需要整合从局部特征到全局依赖的多尺度信息。然而,现有方法在建模长距离全局信息方面存在挑战,其中卷积神经网络受到其局部感受野的限制,视觉变换器因其注意力机制的高二次复杂性而受到限制。最近,基于Mamba的模型因其在长序列建模方面的卓越能力而受到广泛关注。一些研究表明,这些模型可以在各种任务中胜过流行的视觉模型,提供更高的准确性、更低的内存消耗和更小的计算负担。然而,现有的基于Mamba的模型大多是从头开始训练的,并没有探索预训练的力量,这在数据高效的医学图像分析中被证明是非常有效的。本文介绍了一种新的基于Mamba的模型Swin-UMamba,专门为医学图像分割任务设计,利用基于ImageNet的预训练的优势。我们的实验结果揭示了基于ImageNet的训练在提高基于Mamba的模型性能中的关键作用。与CNN、ViT和最新的基于Mamba的模型相比,Swin-UMamba表现出卓越的性能。值得注意的是,在AbdomenMRI、Ecoscopy和Microscopy数据集上,Swin-UMamba的平均得分比其最接近的对手U-Mamba高出2.72%。

关键词

医学图像分割 · 基于ImageNet的预训练 · 长距离依赖建模

2 方法

我们在图1中展示了Swin-UMamba的整体架构。它主要由1)一个基于Mamba的编码器组成,该编码器在大规模数据集(即ImageNet)上进行了预训练,以提取不同尺度的特征,2)一个解码器,有几个上采样块用于预测分割结果,以及3)跳跃连接,以弥合低级细节和高级语义之间的差距。我们将在以下部分介绍Swin-UMamba的详细结构。

2.1 基于Mamba的VSS块

Mamba[4]使用空间状态序列模型(SSMs)[5]将注意力的复杂性从二次降低为线性,以实现自然语言处理中长序列建模。然而,2D视觉数据和1D语言序列之间的区别需要在将Mamba适应于视觉任务时进行仔细考虑。例如,虽然2D空间信息在视觉任务中至关重要[16],但它不是1D序列建模的主要焦点。直接将Mamba应用于扁平化图像将不可避免地导致受限的感受野,其中无法估计与未扫描补丁的关系。基于[16]的见解,我们在Swin-UMamba中引入了视觉状态空间(VSS)块作为基本单元。VSS块通过使用基于选择性扫描空间状态序列模型(S6)的2D选择性扫描(SS2D)来解决与2D图像数据相关的挑战。给定输入特征,SS2D的输出特征可以写成:
其中是四种不同的扫描方向。expand(·)和merge(·)分别对应于[16]中的扫描扩展和扫描合并操作。S6允许1D数组中的每个元素(例如,文本序列)通过压缩的隐藏状态与任何先前扫描的样本进行交互。我们参考[16]以获取有关S6的更多详细信息。VSS块的整体结构在图1中进行了说明。

2.2 整合基于ImageNet的预训练

主要挑战在于有效地将通用预训练模型整合到分割任务中。为此,我们构建了一个编码器,其结构与在广泛的ImageNet数据集上预训练的VMamba-Tiny[16]相似。它允许我们整合通用视觉模型的力量,以提取具有长距离建模能力的信息,模仿过拟合的风险,并为Swin-UMamba建立一个稳健的初始化。如图1所示,Swin-UMamba的编码器可以分为5个阶段。第一阶段是一个用于2×下采样的卷积层。它与VMamba不同,因为我们更喜欢一个渐进的下采样过程来保留低级细节,这对于医学图像分割很重要[23,25]。后续阶段遵循VMamba-Tiny的设计,每个阶段由一个用于2×下采样的补丁合并层和几个VSS块组成,用于处理高级特征。具体来说,第二阶段的补丁合并层被一个2×2补丁嵌入层替换。每个阶段的VSS块和特征维度分别为{0, 2, 2, 9, 2}和D = {48, 96, 192, 384, 768}。我们用ImageNet预训练的VMamba-Tiny初始化VSS块和补丁合并层。

2.3 Swin-UMamba解码器

Swin-UMamba的解码器遵循常用的U形架构。如图1所示,Swin-UMamba使用一个上采样块,其中包括1)一个带有残差连接的额外卷积块来处理跳跃连接特征,以及2)在每个尺度上增加一个用于深度监督[13]的额外分割头。给定跳跃连接特征和来自上一个上采样块的特征,阶段-l的输出特征和分割图可以表示如下:
其中Cat(·)、DeConv_l(·)、Conv_l(·)分别是特征连接操作、转置卷积和1×1卷积。Res(1)_l(·)和Res(2)_l(·)是阶段-l的两个带有残差连接的卷积块,每个Res(·)由两个带有LeakyRELU激活的卷积层组成。我们使用1×1卷积将特征图维度dl投影到类别数K,以获得最终的分割输出。

2.4 Swin-UMamba†:带有Mamba基解码器的Swin-UMamba

为了进一步探索Mamba在医学语义分割中的潜力,我们提出了一个变体Swin-UMamba†,它带有一个基于Mamba的解码器,可以在大大降低复杂性的同时展现出不错的性能。在Swin-UMamba†上进行了几项修改。首先,上采样块被2×补丁扩展层[3]和两个VSS块替换。我们发现,许多参数和计算负担是由重量级的CNN基解码器引起的。其次,我们将编码器改回VMamba的原始设计,并移除了相应的跳跃连接和多余的上采样块。在解码器中的最后一个补丁扩展层是4×上采样操作,与4×补丁嵌入层相镜像。由于在1×2尺度上没有特征,因此在分辨率{1×, 1/4×, 1/8×, 1/16×}上应用了深度监督。结合所有这些修改,网络参数的数量从60M减少到27M,并且在AbdomenMRI数据集上的FLOPs从68.0G减少到18.9G。有关Swin-UMamba†的更多详细信息,请参阅补充材料。

3 实验

声明

本文内容为论文学习收获分享,受限于知识能力,本文对原文的理解可能存在偏差,最终内容以原论文为准。本文信息旨在传播和学术交流,其内容由作者负责,不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题,请及时与我们联系,我们将在第一时间回复并处理。

下载1:Pytoch常用函数手册
在「深度学习爱好者」公众号后台回复:Pytorch常用函数手册,即可下载全网第一份Pytorch常用函数手册,涵盖Tensors介绍、基础函数介绍、数据处理函数、优化函数、CUDA编程、多线程处理等十四章章内容。
下载2:Python视觉实战项目52讲
小白学视觉公众号后台回复:Python视觉实战项目即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目,助力快速学校计算机视觉。

交流群


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

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