Swin-UMamba：基于Mamba和ImageNet预训练的医学图像分割模型本文提出了Swin-UMamba，借助I

前言本文提出了Swin-UMamba，借助ImageNet预训练模型的力量进一步提升基于Mamba的模型在医学图像分割任务中的性能。

本文转载自CVer

仅用于学术分享，若侵权请联系删除

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

CV各大方向专栏与各个部署框架最全教程整理

【CV技术指南】CV全栈指导班、基础入门班、论文指导班全面上线!!

代码&模型：github.com/JiarunLiu/S…

论文：arxiv.org/abs/2402.03…

Highlight

相比较于U-Mamba，Swin-UMamba在三种医学图像分割任务指标上可以取得平均3.58%的提升。
该篇工作通过实验验证了ImageNet预训练对基于Mamba的医学图像分割模型起到非常重要的作用，在迭代次数不变的情况下最高可为Swin-UMamba带来13.08%的DSC提升。
提出了一种变体网络Swin-UMamba，其仅需要相比于U-Mamba不到1/2的网络参数量和约1/3的FLOPs就能够实现与Swin-UMamba相近的性能。

研究背景介绍

对精准的医学图像分割来说，局部信息建模和全局信息建模都很重要，目前两种主流的网络结构CNNs和ViTs在全局信息建模上都有一定的缺陷：CNNs在局部特征提取上非常有效，但由于其固有的局部归纳偏置导致感受野受限；ViTs具备处理全局上下文信息的能力，然而ViT的注意力机制具有较高的二次计算复杂度，这不利于高分辨率图像的建模，高分辨率图像在医学领域并不罕见。

最近的一些工作尝试将Mamba应用到视觉任务中，例如U-Mamba、VMamba和Vim。基于Mamba的模型优势十分鲜明：准确率更高、显存占用更少、计算开销更低。然而，这些工作主要关注于视觉模型结构设计，模型本身往往是从零开始训练的，没有进一步探讨预训练Mamba模型在医学图像分割中的作用——预训练已经被证明在CNNs和ViTs中是非常有效的，尤其是考虑到许多医学图像数据集数据量较少，使用预训练模型可以有效缓解模型过拟合问题并提升模型泛化能力。因此该工作提出了Swin-UMamba，希望借助ImageNet预训练模型的力量进一步提升基于Mamba的模型在医学图像分割任务中的性能。

方法

如上图所示，Swin-UMamba使用类似于UNet的结构，包括编码器、解码器以及两者之间的跳跃连接。作者在设计Swin-UMamba时首先考虑的问题是如何更好地利用ImageNet预训练模型学习到的多尺度特征进行医学图像分割任务，为此Swin-UMamba的编码器部分基本遵循VMamba-Tiny的网络结构设计以便加载预训练模型参数，主要由patch merging模块和VSS模块构成。与VMamba不同的是，Swin-UMamba的第一个stage使用一个卷积层进行2倍下采样，这是为了使每一个stage的下采样倍数都能保持一致，对应地，Swin-UMamba将patch embedding的图像块尺寸也由调整为。对于解码器部分，Swin-UMamba在上采样模块中加入了一个额外的残差卷积模块用来处理跳跃连接的特征，并在每个尺度上都加入了一个额外的分割头用于深度监督，上采样模块如图1右上角所示。

为了进一步探索Mamba在医学图像分割中的潜力，并验证预训练模型在不同网络结构下的影响，作者还提出了一种变体网络Swin-UMamba 。相比于Swin-UMamba，Swin-UMamba 的网络参数更少、FLOPs更低，且仍能保持与Swin-UMamba相近的结果。为此做了几点修改：首先上采样模块被替换为patch expanding和2个VSS Block，这可以有效地降低网络参数量和计算开销；其次，Swin-UMamba解码器的最后使用了一个4倍上采样的patch expanding操作；对应地，作者去除了Swin-UMamba stage 1中的卷积层，并将编码器中patch embedding的图像块尺寸恢复为。通过上述操作，Swin-UMamba的参数量从40M降至27M、FLOPs从58.4G降至15G，Swin-UMamba的网络结构如图2所示。

实验结果

实验结果如表1、表2和表3所示，在AbdomenMRI（腹部MRI器官分割）、Endoscopy（内窥镜手术器械分割）和Microscopy（细胞分割）三个数据集上，Swin-UMamba和Swin-UMamba相比于baseline方法均有显著的提升。例如，Swin-UMamba相比较于U-Mamba_Enc的平均指标提升3.8%，Swin-UMamba相较于U-Mamba_Enc的平均指标提升2.4%。这里需要注意的一点是Swin-UMamba仅需U-Mamba_Enc不到1/2的网络参数量和约1/3的FLOPs，因此Swin-UMamba的性能提升也是非常可观的。

作者发现ImageNet预训练对Swin-UMamba和Swin-UMamba都起到了非常重要的作用，包括更高的分割精度、更稳定的收敛、减轻过拟合问题、数据效率和较低的计算开销。例如，相较于不使用ImageNet预训练，使用ImageNet预训练能够为Swin-UMamba在Endoscopy上的DSC带来13.08%的显著提升，这可能是因为Endoscopy数据集较小，导致模型更容易过拟合。此外，Swin-UMamba仅需baseline算法1/10的训练迭代次数就能够在AbdomenMRI数据上收敛，这有利于节省训练期间的计算开销。实验中还观察到了一个非常有趣的现象，ImageNet预训练有时对模型收敛稳定性起到至关重要的作用，Swin-UMamba在没有使用预训练模型时很难在AbdomenMRI数据上正常收敛，而当使用预训练模型后，Swin-UMamba又能够以较少的参数量和FLOPs超过所有baseline算法，这也进一步佐证了预训练对Mamba-based模型的重要性。

最后提供一些可视化案例供大家参考：

有关Swin-UMamba的更多细节可以参考arXiv上的文章，模型和代码都已在GitHub上开源，欢迎大家一起交流讨论。

Reference

J. Liu et al., “Swin-UMamba: Mamba-based UNet with ImageNet-based pretraining.” arXiv, Feb. 05, 2024.
Y. Liu et al., “VMamba: Visual State Space Model.” arXiv, Jan. 18, 2024.
J. Ma, F. Li, and B. Wang, “U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation.” arXiv, Jan. 09, 2024.
L. Zhu, B. Liao, Q. Zhang, X. Wang, W. Liu, and X. Wang, “Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model.” arXiv, Jan. 17, 2024.

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

计算机视觉入门1v3辅导班

【技术文档】《从零搭建pytorch模型教程》122页PDF下载

QQ交流群：470899183。群内有大佬负责解答大家的日常学习、科研、代码问题。

其它文章

分享一个CV知识库，上千篇文章、专栏，CV所有资料都在这了

明年毕业，还不知道怎么做毕设的请抓紧机会了

LSKA注意力 | 重新思考和设计大卷积核注意力，性能优于ConvNeXt、SWin、RepLKNet以及VAN

CVPR 2023 | TinyMIM：微软亚洲研究院用知识蒸馏改进小型ViT

ICCV2023|涨点神器！目标检测蒸馏学习新方法，浙大、海康威视等提出

ICCV 2023 Oral | 突破性图像融合与分割研究：全时多模态基准与多交互特征学习

听我说，Transformer它就是个支持向量机