FastSAM3D: An Efficient Segment Anything Model for 3D Volumetric Medical Images

📄 arXiv: 2403.09827v1 📥 PDF

作者: Yiqing Shen, Jingxing Li, Xinyuan Shao, Blanca Inigo Romillo, Ankush Jindal, David Dreizin, Mathias Unberath

分类: eess.IV, cs.CV

发布日期: 2024-03-14

🔗 代码/项目: GITHUB

💡 一句话要点

提出FastSAM3D以解决3D医学图像分割的高延迟问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱三：空间感知与语义 (Perception & Semantics)

关键词: 3D医学图像 分割模型 快速推理 层级蒸馏 稀疏注意力 深度学习 计算机视觉

📋 核心要点

  1. 现有的SAM在3D医学图像分割中面临高内存需求和长处理延迟的问题，限制了其交互使用。
  2. FastSAM3D通过层级渐进蒸馏和3D稀疏闪存注意力机制，显著加速了推理速度，降低了内存需求。
  3. 实验结果显示，FastSAM3D在同样的3D体积上，相比于2D SAMs和3D SAMs分别实现了527.38倍和8.75倍的速度提升。

📝 摘要（中文）

Segment Anything Models (SAMs)因其在未见类别和领域的零-shot泛化能力而受到关注。交互性是SAMs的一个关键优势，用户可以迭代提供提示以细化输出。然而，3D医学成像任务的交互使用需要快速推理时间。现有方法在处理3D体积时面临高内存需求和长处理延迟的限制。为了解决这些挑战，本文提出了FastSAM3D，能够在NVIDIA A100 GPU上以每128128128的3D体积图像8毫秒的速度加速SAM推理。实验表明，FastSAM3D在不显著降低性能的情况下，相比于2D SAMs实现了527.38倍的速度提升。

🔬 方法详解

问题定义：本文旨在解决现有SAM在3D医学图像分割中存在的高延迟和高内存需求问题。现有的2D SAM在处理3D体积时需要重复计算，而3D SAM则面临模型参数和计算量的指数级增长。

核心思路：FastSAM3D的核心思路是通过层级渐进蒸馏技术将复杂的12层ViT-B模型的知识转移到轻量级的6层ViT-Tiny模型中，同时引入3D稀疏闪存注意力机制，以减少内存需求并提高并行处理能力。

技术框架：FastSAM3D的整体架构包括两个主要模块：首先是层级渐进蒸馏模块，用于知识转移；其次是3D稀疏闪存注意力模块，用于优化注意力计算。

关键创新：最重要的技术创新在于结合了层级渐进蒸馏和3D稀疏闪存注意力机制，这与传统的SAM方法相比，显著减少了内存使用并加快了推理速度。

关键设计：在参数设置上，FastSAM3D使用了6层ViT-Tiny作为编码器，采用了新的损失函数以优化蒸馏过程，并在注意力机制中引入了稀疏性设计，以提高计算效率。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，FastSAM3D在处理3D体积图像时，推理速度达到了每128128128图像8毫秒，相比于2D SAMs提升了527.38倍，较3D SAMs提升了8.75倍，且在性能上没有显著下降，展示了其卓越的效率。

🎯 应用场景

FastSAM3D在3D医学图像分割领域具有广泛的应用潜力，能够支持临床医生在实时环境中进行高效的图像分析和决策。其低成本的交互式分割能力将推动医学影像处理技术的普及与发展，未来可能在手术规划、疾病诊断等方面发挥重要作用。

📄 摘要（原文）

Segment anything models (SAMs) are gaining attention for their zero-shot generalization capability in segmenting objects of unseen classes and in unseen domains when properly prompted. Interactivity is a key strength of SAMs, allowing users to iteratively provide prompts that specify objects of interest to refine outputs. However, to realize the interactive use of SAMs for 3D medical imaging tasks, rapid inference times are necessary. High memory requirements and long processing delays remain constraints that hinder the adoption of SAMs for this purpose. Specifically, while 2D SAMs applied to 3D volumes contend with repetitive computation to process all slices independently, 3D SAMs suffer from an exponential increase in model parameters and FLOPS. To address these challenges, we present FastSAM3D which accelerates SAM inference to 8 milliseconds per 128128128 3D volumetric image on an NVIDIA A100 GPU. This speedup is accomplished through 1) a novel layer-wise progressive distillation scheme that enables knowledge transfer from a complex 12-layer ViT-B to a lightweight 6-layer ViT-Tiny variant encoder without training from scratch; and 2) a novel 3D sparse flash attention to replace vanilla attention operators, substantially reducing memory needs and improving parallelization. Experiments on three diverse datasets reveal that FastSAM3D achieves a remarkable speedup of 527.38x compared to 2D SAMs and 8.75x compared to 3D SAMs on the same volumes without significant performance decline. Thus, FastSAM3D opens the door for low-cost truly interactive SAM-based 3D medical imaging segmentation with commonly used GPU hardware. Code is available at https://github.com/arcadelab/FastSAM3D.