人工智能咨询培训老师叶梓 转载标明出处
在医学图像分割的研究中,卷积神经网络(CNN)与变换器(Transformer)模型各有其优势和局限。CNN由于其有限的局部感受野,在捕获长距离信息时能力较弱;而尽管Transformer在全局信息建模上具有显著优势,但其自注意力机制随着图像尺寸的增加,计算复杂度会以二次方的速度增长,这使得其在处理大规模图像数据时面临较大的计算压力。针对这些问题,上海交通大学的研究者提出了一种基于状态空间模型(State Space Models, SSMs)的新型U型网络架构——Vision Mamba UNet(VM-UNet)。该模型不仅在长距离依赖关系的建模上表现出色,而且保持了线性的计算复杂度。
方法
图 1(a) 展示了VM-UNet的整体架构。VM-UNet包括一个Patch Embedding层、一个编码器、一个解码器、一个Final Projection层和跳跃连接。与之前的方法不同,采用了非对称设计。
Patch Embedding层
该层将输入图像x划分为不重叠的4×4大小的块,然后将图像的维度映射到C维,其中C默认为96。这个过程产生了嵌入后的图像x',尺寸为R H4 × W4 × C。在将x'送入编码器进行特征提取之前,使用Layer Normalization对其进行了归一化。
编码器
编码器由四个阶段组成,在前三个阶段的末尾应用了patch merging操作,以减少输入特征的高和宽,同时增加通道数。在四个阶段中分别使用了**[2, 2, 2, 2] VSS块**,每个阶段的通道数分别为[C, 2C, 4C, 8C]**。
解码器
解码器同样分为四个阶段。在最后三个阶段的开始,使用patch expanding操作来减少特征通道的数量并增加高度和宽度。在四个阶段中,使用了**[2, 2, 2, 1] VSS块**,每个阶段的通道数分别为[8C, 4C, 2C, C]。
Final Projection层
解码器之后是一个Final Projection层,用于将特征的大小恢复到与分割目标匹配。具体来说,通过patch expanding进行4倍上采样以恢复特征的高和宽,然后通过一个投影层来恢复通道数。
跳跃连接
对于跳跃连接,采用了简单的加法操作,没有引入任何额外的参数。
VSS块是VM-UNet的核心模块,源自VMamaba。在经过Layer Normalization后,输入被分成两个分支。第一个分支通过一个线性层后接一个激活函数;第二个分支则通过一个线性层、一个深度可分离卷积和一个激活函数处理后,送入2D-Selective-Scan (SS2D)模块进行进一步的特征提取。然后,使用Layer Normalization对特征进行归一化,接着执行元素间乘法操作,将两个分支的输出合并。最后,使用一个线性层混合特征,并通过与残差连接的组合形成VSS块的输出。默认使用SiLU作为激活函数。
SS2D由三个组件组成:scan expanding操作、S6块和scan merging操作。如图2(a)所示,scan expanding操作将输入图像沿四个方向(从左上到右下、从右下到左上、从右上到左下、从左下到右上)展开成序列。然后,这些序列由S6块处理以提取特征,确保从各个方向的信息都被彻底扫描,从而捕获多样化的特征。如图2(b)所示,scan merging操作将四个方向的序列求和并合并,恢复输出图像到与输入相同的尺寸。S6块在S4的基础上引入了选择性机制,通过根据输入调整SSM的参数,使模型能够区分并保留相关信息,同时过滤掉不相关的信息。算法 1展示了S6块的伪代码。
VM-UNet的引入旨在验证基于纯SSM模型在医学图像分割任务中的应用潜力。因此,专门使用了最基本的二元交叉熵和Dice损失(BceDice loss)以及交叉熵和Dice损失(CeDice loss)作为二元和多类分割任务的损失函数。公式5和6分别表示这两种损失函数:
其中,N表示样本总数,C代表类别总数。y_i和ŷ_i分别表示真实标签和预测。y_i,c是一个指示器,如果样本i属于类别c,则为1,否则为0。ŷ_i,c是模型预测样本i属于类别c的概率。|X|和|Y|分别代表真实值和预测值。λ_1和λ_2指的是损失函数的权重,这里默认都设置为1。
想要掌握如何将大模型的力量发挥到极致吗?叶老师带您深入了解 Llama Factory —— 一款革命性的大模型微调工具。实战专家1小时讲解让您轻松上手,学习如何使用 Llama Factory 微调模型。
评论留言“参加”或扫描微信备注“参加”,即可参加线上直播分享,叶老师亲自指导,互动沟通,全面掌握Llama Factory。关注享粉丝福利,限时免费录播讲解。
实验
实验部分在ISIC17、ISIC18和Synapse数据集上评估了VM-UNet在医学图像分割任务上的性能。
ISIC17和ISIC18数据集:国际皮肤成像协作2017年和2018年挑战数据集(ISIC17和ISIC18)包含2,150张和2,694张皮肤病变分割数据集,分别带有分割掩码标签。按照之前的工作,数据集按照7:3的比例划分为训练集和测试集。对于ISIC17数据集,训练集包括1,500张图像,测试集包括650张图像。对于ISIC18数据集,训练集包括1,886张图像,测试集包含808张图像。对于这两个数据集,提供了包括平均交并比(mIoU)、Dice相似系数(DSC)、准确率(Acc)、敏感性(Sen)和特异性(Spe)在内的多个指标的详细评估。
Synapse多器官分割数据集:Synapse数据集包含30个腹部CT病例,共3,779张腹部临床CT图像,涵盖8种腹部器官(主动脉、胆囊、左肾、右肾、肝脏、胰腺、脾脏、胃)。按照之前工作的设置,使用18个病例进行训练,12个病例进行测试。对于这个数据集,报告了Dice相似系数(DSC)和95% Hausdorff距离(HD95)作为评估指标。
ISIC17和ISIC18数据集中的图像调整为256×256大小,而Synapse数据集中的图像调整为224×224大小。为了防止过拟合,采用了包括随机翻转和随机旋转在内的数据增强技术。对于ISIC17和ISIC18数据集,使用了BceDice损失函数,而对于Synapse数据集,则采用了CeDice损失函数。批量大小设置为32,并采用AdamW优化器,初始学习率设置为1e-3。使用CosineAnnealingLR作为调度器,最大迭代次数设为50次,最小学习率设为1e-5。训练周期设置为300。对于VM-UNet,使用在ImageNet-1k上预训练的VMamba-S的权重初始化编码器和解码器。所有实验都在单个NVIDIA RTX A6000 GPU上进行。
表1展示了ISIC17和ISIC18数据集上的比较实验结果。表2展示了Synapse数据集上的比较实验结果,其中也报告了每个类别的DSC。对于ISIC17和ISIC18数据集,VM-UNet在mIoU、DSC和Acc指标上优于其他模型。对于Synapse数据集,VM-UNet也取得了有竞争力的性能。例如,在DSC和HD95指标上,该模型超过了Swin-UNet(首个纯Transformer基模型)1.95%和2.34mm。结果证明了基于SSM模型在医学图像分割任务上的优越性。
使用ISIC17和ISIC18数据集对VM-UNet的初始化进行了消融实验。分别使用VMamba-T和VMamba-S的预训练权重初始化VM-UNet。表3中的实验结果表明,更强大的预训练权重显著提高了VM-UNet的下游性能,表明VM-UNet在很大程度上受到预训练权重的影响。
实验结果显示,VM-UNet在多个评价指标上均优于其他模型,包括平均交并比(mIoU)、Dice相似系数(DSC)和准确率(Acc)。
代码已经开源,感兴趣的读者可以通过以下链接访问项目: