Swin-UMamba：基于Mamba的UNet与基于ImageNet的预训练
图像与视频处理
2024-03-07 v2 计算机视觉与模式识别
机器学习
摘要
精确的医学图像分割需要整合从局部特征到全局依赖的多尺度信息。然而，现有方法难以建模长程全局信息，其中卷积神经网络（CNNs）受限于其局部感受野，而视觉Transformer（ViTs）则因其注意力机制的高二次复杂度而困扰。最近，基于Mamba的模型因其在长序列建模方面的出色能力而备受关注。多项研究表明，这些模型在各种任务中能超越流行的视觉模型，提供更高的精度、更低的内存消耗和更少的计算负担。然而，现有的基于Mamba的模型大多是从头开始训练的，并未探索预训练的强大能力，而预训练已被证明对于数据高效的医学图像分析非常有效。本文介绍了一种新颖的基于Mamba的模型Swin-UMamba，专为医学图像分割任务设计，利用了基于ImageNet的预训练优势。我们的实验结果揭示了基于ImageNet的训练在提升基于Mamba模型性能方面的关键作用。与CNNs、ViTs和最新的基于Mamba的模型相比，Swin-UMamba展现出了大幅领先的优越性能。值得注意的是，在AbdomenMRI、内窥镜和显微镜数据集上，Swin-UMamba的平均得分比其最接近的对应模型U-Mamba_Enc高出2.72%。
引用
@article{arxiv.2402.03302,
title = {Swin-UMamba: Mamba-based UNet with ImageNet-based pretraining},
author = {Jiarun Liu and Hao Yang and Hong-Yu Zhou and Yan Xi and Lequan Yu and Yizhou Yu and Yong Liang and Guangming Shi and Shaoting Zhang and Hairong Zheng and Shanshan Wang},
journal= {arXiv preprint arXiv:2402.03302},
year = {2024}
}
备注
Code and models of Swin-UMamba are publicly available at: https://github.com/JiarunLiu/Swin-UMamba