0-MedficientSAM: A Robust Medical Segmentation Model with Optimized Inference Pipeline for Limited Clinical Settings
MedficientsAM:具有优化推理管道的强大医学分割模型,用于有限的临床环境
背景
SAM的出现使通用医学图像分割模型的开发跨越了不同的模式。但是,临床环境中这种深度学习模型的推理仍然受到对强大计算设备的依赖的限制。
医学图像的早期分割模型通常基于nnUNET结构。尽管有效,但这些模型仅限于特定数据集,每个模型都针对特定的分割任务量身定制。SAM的出现标志着分割任务的显着范式移动。这些努力的重点是将SAM适应医疗数据,同时保持较高的计算需求。在大多数医疗机构中,不可用强大的计算设备,需要快速的结果,因此在实践中部署这些模型的挑战。
因此本文提出MedficientSAM,通过将EfficientViT模型与医学图像分割任务深度融合,提出了一种高效、精准且可临床落地的通用医学分割解决方案。其核心贡献在于:首次将轻量级ReLU线性注意力引入医学领域,通过两阶段训练(知识蒸馏+端到端调优)从MedSAM迁移知识至EfficientViT编码器,在保留SAM强泛化能力的同时,将模型计算复杂度降低至原版的1/6;创新性地采用C++重构推理管线并设计嵌入缓存机制,使CPU端2D/3D医学影像分割速度分别提升5倍和4.6倍(相比Python版LiteMedSAM),内存占用减少50%;在公开验证集上以86.42% DSC和87.95% NSD刷新性能记录,并在”Segment Anything In Medical Images On Laptop”挑战赛中综合排名第一,为CT、MRI、超声等多模态医学影像在边缘设备的实时分析提供了首个兼顾效率与精度的通用分割框架,显著推进了AI辅助诊断的临床实用性。
实验方法
MedficientSAM 基于 EfficientViT-SAM。用 EfficientViT 替换了 MedSAM 的图像编码器。像 EfficientViT-SAM 一样,MedficientSAM 有三种变体,分别是 L0、L1 和 L2,按模型大小递增顺序列出。该模型分两个阶段进行训练:蒸馏和微调。
EfficientViT-SAM-L1的宏观体系结构(上)和MedficientSAM(下)
数据预处理
对齐所使用的全部原图和MASK大小,原图使用双线性插值进行调整,应用最小-最大缩放,并用零值填充调整后的图像以创建正方形尺寸。MASK使用最近邻插值,然后用零填充调整后的掩模以达到正方形尺寸
蒸馏
为了初始化图像编码器,通过蒸馏将 MedSAM 的图像编码器(ViT-B)的知识转移到 EfficientViT。目标是通过最小化 L2 损失函数,使 EfficientViT 和 MedSAM-ViT-B 的图像嵌入对齐。
微调
将蒸馏的 EfficientViT 与 MedSAM 的预训练提示编码器和掩码解码器集成,创建了 MedficientSAM。随后,进行端到端的训练,以进一步提高性能。为了提示模型,通过确定覆盖二进制掩码的最小矩形生成框提示,并引入随机偏移以提高模型的鲁棒性。
损失函数
使用Focal损失和Dice损失的加权求和,比例为20:1。
3D推理
受到LiteMedSAM的启发,在进行3D体积推断时,从中间切片开始,向两端传播,使用之前预测的掩膜切片作为指导提示。采用了在之前的器官分割工作中使用的掩膜传播的理念。如果在之前的切片中找到二进制掩膜,会获取覆盖该二进制掩膜的边界框,并将其作为当前切片的框提示,而不是使用输入中的框提示。
后处理
MedficientSAM输出的二进制掩码的固定尺寸为256 × 256。首先调整这些输出掩码的大小以匹配图像编码器的输入大小,然后将填充零件裁剪出来,并最终将其调整为原始分辨率
推理优化
尽管对于模型原型制作非常方便,但Python由于其解释性质而不适合部署。所以作者将管道移植到C++,并使用OpenVino作为模型运行时来减少推理时间。
实验结果
分割涵盖了11种医学图像方式(CT、MRI、PET、X 光、超声、乳腺摄影、OCT、内窥镜、眼底镜、皮肤镜和显微镜)。评估指标包括两种准确度测量:Dice 相似系数 (DSC) 和标准化表面 Dice (NSD),还有运行时间作为效率测量。
定量结果:
公共验证集的定量评估结果。
定性结果:
来自各种公共数据集的定性结果。
总结
这项工作介绍了MedficientSAM,利用EfficientViT来提高MedSAM的效率和准确性。方法采用两阶段的训练过程,与MedSAM相比,分割准确性得到了提高,同时显著降低了计算需求。此外,开发了一种新颖的C++推理管道,使MedficientSAM能够在临床环境中常见的资源受限设备上运行。