本文主要是介绍ViM-UNet:用于生物医学细分的 Vision Mamba,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!
ViM-UNet:用于生物医学细分的 Vision Mamba
- 摘要
- Introduction
- Method and Experiments
- 结果与讨论
ViM-UNet: Vision Mamba for Biomedical Segmentation
摘要
卷积神经网络(CNNs),尤其是UNet,是生物医学分割的默认架构。已经提出了基于Transformer的方法,如UNETR,以取代它们,得益于全局视野,但受到更大的运行时间和更高的参数数量的影响。
最近的Vision Mamba架构为Transformer提供了一个有吸引力的替代方案,同样提供了全局视野,但效率更高。
在这里,作者介绍了基于它的ViM-UNet,一种新颖的分割架构,并将其与UNet和UNETR在两个具有挑战性的显微实例分割任务上进行比较。作者发现,它在不同任务上表现得与UNet相似或更好,并且在效率更高的情况下优于UNETR。
代码地址:
https://github.com/constantinpape/torch-em/blob/main/vimunet.md
Introduction
分割是生物医学图像分析中的一个重要任务,应用范围从放射学到显微镜学。大多数现代分割方法都基于卷积神经网络(CNN),其中UNet(Ronneberger等人,2015年)最受欢迎。在文本和视觉领域(ViT,(Dosovitskiy等人,2021年))的成功之后, Transformer 架构也被提出来用于生物医学分割;尤其是UNETR(Hatamizadeh等人,2021年)和SwinUNETR(Hatamizadeh等人,2022年)。它们具有全局视野,对于需要大上下文的任务来说有望获得更好的质量。然而,它们的运行时间和参数数量也更大。最近,Mamba架构(Gu和Dao,2023年)被提出,该架构扩展了状态空间模型(SSM)(Gu等人,2022年),旨在克服这些计算效率低的问题,同时保持全局视野。它已经被Vision Mamba(ViM)(Zhu等人,2024年)适应用于计算机视觉。
在这里,作者介绍了基于 ViM的 ViM-UNet,用于生物医学分割,并将其与 UNet 和 UNETR进行了比较,用于显微镜实例分割,这是生物学的一个重要分析任务。这项任务的大多数方法,例如CellPose(Stringer等人,2021年),StarDist(Schmidt等人,2018年),都是基于UNet架构,而最近的方法也采用了 Transformer ,例如(Archit等人,2023年),并且当前的基准测试(Ma等人,2024年)显示了 Transformer 有利的成果。作者使用了两个具有不同特点的数据集,见图1,发现ViM-UNet的性能与UNet相当或更好(取决于任务),而UNETR表现不佳。作者对结果进行了外部方法的验证,nnUNet(Isensee等人,2021年),这是一个经过良好测试的UNet框架,以及U-Mamba(Ma等人,2024年),它也基于Mamba,但缺少ViM的视觉特定优化。作者的结果显示了ViM在生物医学图像分析中的潜力。作者认为它特别有希望用于依赖于大上下文的任务,例如3D分割或细胞追踪。
Method and Experiments
作者比较了三种不同的架构:
UNet、UNETR和作者贡献的ViM-UNet,这些架构都在torch-em(Pape)中实现。UNet有4个层级,初始特征为64,每提升一个层级特征数量翻倍。对于UNETR,作者使用了Segment Anything(Kirillov等人,2023年)的ViT和与UNet相同的解码器,每个层级由两个卷积层和一个转置卷积层组成;输入来自前一个解码器层和ViT的输出。
作者选择了这种简单的实现方式,而不是像UNet和原始UNETR中的跳跃连接,后者将编码器和解码器中的相应层级连接起来。
作者发现这种设计没有负面影响。对于ViM-UNet,作者使用带有双向SSM层的ViM编码器。与ViT类似,这个模型在图像块上操作;作者使用的图像块大小为16。解码器设计与UNETR相同。对于UNETR和ViM-UNet,作者比较了不同大小的编码器;ViT有Base、Large、Huge三种尺寸,而ViM有Tiny、Small两种尺寸。
图1:使用ViM-UNetSmall进行分割的示例图像。
作者对两个数据集进行了比较:相位对比显微镜下的细胞分割(LIVECell (Edlund et al., 2021))以及体积电子显微镜下的神经纤维分割(CREMI (Funke et al., 2016))。LIVECell 包含了形态各异的微小细胞,而CREMI 包含了大小不一的神经纤维,详见图1。作者将CREMI的分割限制为2D。对于LIVECell,作者使用了给定的训练、验证和测试划分;对于CREMI,作者使用每个 Voxel 的前75个切片进行训练,接下来的25个进行验证,最后的25个进行测试。在LIVECell实例分割中,作者预测(i)前景和边界概率,这之后通过水波算法进行后处理,以及(ii)前景概率以及细胞中心和边界的距离,同样也用水波算法进行后处理。作者选择(i)与其他实现(见下文)进行比较,其中作者无法实现距离预测,而(ii)因为这种方法更适合这里。对于CREMI作者使用(i)。边界预测是这个任务的常规步骤,通常接下来会进行图聚集(Beier et al., 2017),这在2D中是不需要的。网络使用Adam进行训练,共10万次迭代,初始学习率为,并在平台期进行减少。作者与带有边界分割的nn-UNet (Isensee et al., 2021) 和 U-Mamba (Ma et al., 2024) 进行比较。这两种方法都通过超参数搜索进行配置,作者使用默认设置。作者使用平均分割准确度(Everingham et al., 2010)进行评估。
结果与讨论
图2展示了结果。对于LIVECell数据集,采用距离分割的UNet表现最佳,紧随其后的是ViM-UNet。UNETR的表现显著较差。对于边界分割,UNet明显表现最好,其次是ViM-UNet和外部方法。UNETR再次表现不佳。对于CREMI数据集,ViM-UNet表现最佳,其次是外部方法和UNet,UNETR的结果较弱。作者假设全局视野对于小结构(如LIVECell)并没有带来任何优势,但ViM-UNet可以利用它来处理大结构(如CREMI)。UNETR表现不佳,这很可能是由于参数数量较多(见表1)且缺乏预训练;请注意,通过预训练可以实现更好的性能(Horst等人,2023年)。与外部方法的比较验证了作者的实现并没有表现不佳,但由于训练和推理的差异,无法进行完全客观的比较。作者还研究了推理时间和训练所需的内存,见表1。UNet是最有效的架构,其次是ViM-UNet和UNETR。
表1:作者模型的参数数量、训练所需的VRAM以及每张图像的推理时间(以秒为单位)。
图2:作者的方法和外部方法的对比结果;圆圈突出了三种最佳方法。
总的来说,ViM-UNet在生物医学图像分析方面很有前景。作者认为,它可能取代基于 Transformer 的方法,在需要大上下文的领域中应用,因为它同样具有全局视野,但效率更高。其较低的参数数量使得可以在较小的数据集上进行应用,且无需大量预训练。作者计划将其扩展到3D分割和跟踪,在这些领域大上下文通常至关重要。
这篇关于ViM-UNet:用于生物医学细分的 Vision Mamba的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!