《VM-UNet: Vision Mamba UNet for Medical Image Segmentation》论文阅读及代码复现

论文地址：[2402.02491] VM-UNet: Vision Mamba UNet for Medical Image Segmentation (arxiv.org)

代码地址：JCruan519/VM-UNet: (ARXIV24) This is the official code repository for “VM-UNet: Vision Mamba UNet for Medical Image Segmentation”. (github.com)

摘要：

在医学图像分割领域，基于细胞神经网络和基于变换器的模型都得到了广泛的探索。然而，细胞神经网络在长期建模能力方面表现出局限性，而变压器则受到其二次计算复杂性的阻碍。最近，以曼巴为例的状态空间模型（SSM）已成为一种很有前途的方法。它们不仅擅长对长程相互作用进行建模，而且保持线性计算复杂性。在本文中，利用状态空间模型，我们提出了一种用于医学图像分割的Ushape架构模型，称为Vision Mamba UNet（VM UNet）。具体地，引入视觉状态空间（VSS）块作为基础块来捕获广泛的上下文信息，并构建了一种不对称的编码器-解码器结构。我们在ISIC17、ISIC18和Synapse数据集上进行了全面的实验，结果表明，VM UNet在医学图像分割任务中具有竞争力。据我们所知，这是第一个基于纯SSM模型构建的医学图像分割模型。我们的目标是建立一个基线，并为未来开发更高效、更有效的基于SSM的分割系统提供有价值的见解。

技术现状

近年来，状态空间模型(State Space Models)引起了研究者的极大兴趣。在经典SSM研究的基础上，现代SSM（例如Mamba）不仅建立了长距离依赖关系，而且在输入大小方面表现出线性复杂性。此外，基于SSM的模型在许多领域得到了大量研究，包括语言理解、一般视觉等。特别是，U-Mamba最近推出了一种新的SSM-CNN混合模型，标志着它首次应用于医学图像分割任务。SegMamba在编码器部分结合了SSM，同时在解码器部分仍然使用CNN，这提出了用于3D脑肿瘤分割任务的SSM-CNN混合模型。尽管上述工作已经将SSM用于医学图像分割任务，但纯基于SSM的模型的性能仍有待探索。

文章贡献

1）提出了VM-UNet，这标志着首次探索了纯基于SSM的模型在医学图像分割中的潜在应用。

2） 在三个数据集上进行了综合实验，结果表明VM UNet表现出相当大的竞争力。

3） 在医学图像分割任务中为纯基于SSM的模型建立了基线，提供了有价值的见解，为开发更高效、更有效的基于SSM分割方法铺平了道路。

整体结构

① Vision Mamba UNet (VM-UNet)

VM-UNet包括Patch Embedding layer, an encoder, a decoder, a Final Projection layer, and skip connections。采用了不对称结构设计。Patch Embedding layer将输入图像划分为4*4的不重叠补丁，随后将图像的维度映射到C（默认96），该过程得到嵌入图像H4×W/4×C，然后编码器的4个stage进行特征提取，前3个stage结束时应用补丁合并操作以减少输入特征的高度和宽度，同时增加通信数量。每个stage使用2个VSS块。解码器部分与编码器类似。

② VSS block

输入数据通过层归一化后，被拆分为两个分支。在第一个分支中，输入通过线性层和激活函数；第二个分支通过线性层、深度可分离卷积和激活函数、2D选择性扫描模块2D-Selective-Scan（SS2D）。随后进行层归一化，与第一个分支的输出执行逐元素生成。最后进行一个残差连接操作得到VSS block的输出。结构中的损失函数选择的SiLU

SS2D由3部分组成，a scan expanding operation, an S6 block, and a scan merging operation

scan expanding operation将输入图像沿着四个不同的方向（从左上到右下、从右下到左上、从右上到左下、从左下到右上）展开为序列。然后S6块对这些序列进行处理，以进行特征提取，确保对来自各个方向的信息进行彻底的扫描，从而获取不同的特征。随后，合并序列，将输入图像恢复到与输入相同的大小。

S6块源自Mamba，通过根据输入调整SSM的参数，在S4之上引入了一种选择性机制。这使得模型能够区分和保留相关信息，同时过滤掉不相关的信息。算法1中给出了S6块的伪代码

结论

在本文中，我们首次引入了一种纯的基于SSM的医学图像分割模型，将VM-UNet作为基线。为了利用基于SSM的模型的能力，我们使用VSS块构建VM UNet，并使用预训练的VMamba-S初始化其权重。在皮肤损伤和多器官分割数据集上进行的综合实验表明，纯基于SSM模型在医学图像分割任务中具有很强的竞争力，值得在未来进行深入探索。

代码复现

注意目前只支持linux系统复现

conda create -n vmunet python=3.8
conda activate vmunet
pip install torch==1.13.0 torchvision==0.14.0 torchaudio==0.13.0 --extra-index-url https://download.pytorch.org/whl/cu117
pip install packaging
pip install timm==0.4.12
pip install pytest chardet yacs termcolor
pip install submitit tensorboardX
pip install triton==2.0.0
pip install causal_conv1d==1.0.0  # causal_conv1d-1.0.0+cu118torch1.13cxx11abiFALSE-cp38-cp38-linux_x86_64.whl
pip install mamba_ssm==1.0.1  
# 如果安不上按下面步骤
# 点击报错中链接
# pip install mamba_ssm-1.0.1+cu118torch1.13cxx11abiFALSE-cp38-cp38-linux_x86_64.whl
pip install scikit-learn matplotlib thop h5py SimpleITK scikit-image medpy yacs
#预训练权重下载地址：https://pan.baidu.com/s/144nc45k28NBB_8k9hf_OeQ?pwd=xuns