当前位置:   article > 正文

SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation

segmamba

arxiv' 2024

paper: https://arxiv.org/pdf/2401.13560.pdf

code: https://github.com/ge-xing/SegMamba

Abstract

Transformer体系结构在建模全局关系方面显示出了非凡的能力。然而,它在处理高维医学图像时提出了重大的计算挑战。这阻碍了它在这项任务中的发展和广泛采用。Mamba作为一种状态空间模型(State Space Model, SSM),近年来作为序列建模中一种值得关注的远程依赖关系模型,以其显著的存储效率和计算速度在自然语言处理领域表现优异。受其成功的启发,我们引入了SegMamba,一种新颖的3D医学图像分割Mamba模型,旨在有效地捕获每个尺度下整个体积特征中的远程依赖关系。与基于transformer的方法相比,从状态空间模型的角度来看,SegMamba在整体体积特征建模方面表现出色,即使体积特征的分辨率为64 × 64 × 64(序列长度约为260k),也能保持卓越的处理速度。在BraTS2023数据集上的综合实验证明了该方法的有效性和高效性。可从https://github.com/ge-xing/SegMamba获得SegMamba的代码。

关键词:状态空间模型·曼巴·多序列建模·三维医学图像分割

1 Introduction

扩展模型的接受野是三维医学图像分割的一个关键方面。传统的卷积神经网络(cnn)在从高分辨率三维医学图像中提取全局信息方面不是很有效。因此,建议使用具有大核大小的深度卷积[13]来建模更广泛的特征。3D UX-Net[10]引入了一种新的架构,利用卷积(7 × 7 × 7)块和一个大内核来促进更大的接受域。然而,当像素之间的距离太大时,基于cnn的方法难以建立关系。

近年来,利用自关注模块提取全局信息的transformer架构被广泛应用于三维医学图像分割。TransBTS[19]结合3D-CNN提取局部空间特征,然后应用transformer对高级特征中的全局依赖关系进行建模。UNETR[5]采用Vision Transformer (ViT)[2]作为编码器来学习上下文信息,然后通过多个分辨率的跳过连接与基于cnn的解码器合并。SwinUNETR[4]利用SwinTransformer[12]作为编码器提取多尺度特征。同时设计了一个多尺度解码器,融合了每个编码器阶段的特征,在三维医学图像分割中取得了很好的效果。然而,对于以transformer为基础的方法,典型的3D医学图像的高分辨率会导致显著的计算负担和降低速度性能。

为了克服长序列建模的挑战,Mamba[3]源于状态空间模型(ssms)[8],旨在通过选择机制和硬件感知算法对远程依赖关系进行建模,提高训练和推理的效率。许多研究探索了Mamba在计算机视觉(CV)中的应用。U-Mamba[14]将Mamba层集成到nnunet[7]的编码器中,以增强一般医学图像的分割。同时,Vision Mamba[21]提出了Vim块,该块结合了双向SSM,用于数据依赖的全局视觉上下文建模和位置嵌入,用于位置感知的视觉理解。此外,VMamba[11]设计了一个CSM模块来弥补一维阵列扫描和二维平面遍历之间的差距。在三维医学图像分割中,传统的transformer块在处理大尺寸特征时面临挑战。为了获得更强的视觉理解,有必要对高维特征中的相关性进行建模。受此启发,我们引入了SegMamba,这是一种新颖的框架,将u形结构与Mamba结合起来,在各种尺度上模拟整个体量的全局特征。据我们所知,这是利用Mamba专门用于3D医学图像分割的第一种方法。与传统的基于cnn的和基于transformer的方法相比,SegMamba在体积数据中表现出卓越的远程依赖关系建模能力,同时保持了出色的推理效率。大量的实验证明了该方法的有效性。

2 Method

SegMamba主要由三个部分组成:1)具有多个Mamba块的Mamba编码器,用于提取不同尺度的特征;2)基于卷积层的三维解码器,用于预测分割结果;3)跳过连接将多尺度特征连接到解码器,用于特征重用。图1说明了所提出的SegMamba的概述。我们在本节中进一步描述编码器和解码器的细节。

2.1 Mamba Encoder

全局特征和多尺度特征建模是三维医学图像分割的重要内容。虽然transformer体系结构可以提取全局信息,但当处理过长的特征序列时,它会产生显着的计算负担。为了减少序列长度,UNETR等基于transformer结构的方法直接对输入的三维医学图像进行下采样,分辨率为D*H*W至D/16 ×H/16 ×W/16。然而,这种方法限制了多尺度特征的建模能力,而多尺度特征对于通过解码器预测分割结果至关重要。为了克服这一限制,我们设计了Mamba块,它用更高效的Mamba层代替了transformer体系结构中的自关注模块。这使得多尺度和全局特征建模成为可能,同时在训练和推理过程中保持高效率。

如图1所示,Mamba编码器由一个stem层和多个Mamba块组成。对于stem层,我们采用深度卷积,内核大小为7 × 7 × 7,填充为3 × 3 × 3,步幅为2 × 2 × 2。给定一个三维输入体积I∈RC×D×H×W,其中C表示输入通道数,stem层提取第一个尺度特征z0∈R48× D/2 × H/2 × W/2。然后,z0通过每个Mamba块和相应的下采样层馈送。值得注意的是,SegMamba在Mamba块之前采用了平坦化操作ϕ,将3D特征重塑为1D长序列,从而实现了具有较少感应偏置的高效序列建模。在Mamba块后,通过σ运算将特征恢复到三维形状。这个过程可以定义为:

其中LN为归一化输入特征的层范数,Mamba为Mamba层,MLP为丰富特征表示的多层感知层。

2.2 Decoder

Mamba编码器提取了多尺度特征,在前人的研究基础上,我们利用基于cnn的解码器和跳跃连接形成U-shape网络来预测分割结果。

3 Experiments

3.1 Datasets

BraTS2023 dataset : BraTS2023数据集[15,1,9]共包含1,251个3D脑MRI体积。每个体积包括四种模式(T1, T1Gd, T2, T2- flair)和三个分割目标(WT:整个肿瘤,ET:增强肿瘤,TC:肿瘤核心)。所有数据都被重新采样到相同的间距(1.0,1.0,1.0)。

3.2 Evaluation Metrics

基于不同数据集的特征,我们采用各种指标来比较我们的方法与其他方法的性能。

Dice similarity coefficient (Dice score):Dice分数是一个重叠度量,用来衡量预测和真实之间的重叠百分比。计算公式如下:

式中,X表示语义预测,Y表示GT,|·|表示基数计算操作。

95% Hausdorff distance (HD95): HD95度量标准测量Hausdorff距离的第95个百分位数,为预测和实际情况GT之间的最大距离提供可靠的评估。由下式定义:

式中,d(X, Y)表示X与Y之间的Hausdorff Distance ,maxk95%表示第95百分位处的最大值。

3.3 Comparison Methods

为了进行全面的评估,将SegMamba与其他七种最先进的方法进行了比较,这些方法都涵盖了CNN和transformer架构。

CNN-based methods。 我们将SegMamba与SegresNet[16]、UX-Net[10]和MedNeXt[17]进行了比较。其中,UX-Net和MedNeXt是最新的3D医学图像分割方法。

Transformer-based methods。我们还比较了SegMamba与最新的基于transformer的在医学成像上的分割方法。UNETR[5]、SwinUNETR[4]、SwinUNETR- v2[6]和nnFormer[20]是最著名的基于transformer的三维医学图像分割方法。他们使用Vision Transformer[2]或SwinTransformer[12]结构作为编码器来建模全局特征。

为了确保公平的比较,我们在实验中使用了所有方法的公开代码,并保持所有设置相同。

3.4 Implementation Details

我们的模型在Pytorch 2.0.1-cuda11.7和Monai 1.2.0中实现。在训练过程中,我们对每个数据集使用128 × 128 × 128的随机裁剪大小和每个GPU bs为2的批处理。由于BraTS2023数据集中的每个体积包含4个模态,因此我们在网络的输入处将通道维度中的每个模态连接起来。我们对所有实验使用交叉熵损失和SGD优化器以及多项式学习率调度器(初始学习率为1e-2,衰减为1e-5)。我们对所有数据集运行1000个epoch,并采用以下数据增强:添加亮度,伽马,旋转,缩放,镜像和弹性变形。所有实验都是在一个云计算平台上进行的,该平台配备了4个NVIDIA A100 gpu。我们将测试时间增强(TTA)技术[18](即镜像预测和重叠滑动窗口推断[7])应用于所有三个数据集的推理,并将重叠比设置为0.5。对于每个数据集,我们随机分配70%的3D体积用于训练,10%用于验证,剩下的20%用于测试,确保每个体积在训练、验证和测试集中只出现一次。

3.5 Quantitative Comparison to Previous Methods

BraTS2023 brain tumor segmentation 。BraTS2023数据集的胶质瘤分割结果如表1所示,其中我们使用Dice评分和HD95来评价在WT、TC和ET三个分割目标上的表现。基于cnn的UX-Net方法在比较方法中表现最好,其平均Dice评分为89.69%,平均HD95为4.81。在基于transformer的方法中,swinunetrt - v2也表现出良好的性能,平均Dice得分为89.39%,平均HD95为4.51。相比之下,我们的SegMamba在WT, TC和ET上的Dice得分分别为93.61%,92.65%和87.71%,HD95得分分别为3.37,3.85,3.48,优于所有其他方法。我们的SegMamba的平均Dice得分为91.32%,比排名第二的UX-Net和排名第三的SwinUNETRv2分别提高了1.63%和1.93%。

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/很楠不爱3/article/detail/636313
推荐阅读
相关标签
  

闽ICP备14008679号