[论文解读] EdgeSAM: Prompt-In-the-Loop Distillation for SAM

Chong Zhou, Xiangtai Li+2 more|arXiv (Cornell University)|Dec 11, 2023

Advanced Memory and Neural Computing参考文献 61被引用 11

一句话总结

EdgeSAM 通过提示环路蒸馏，将 SAM 精炼为基于 CNN 的编码器，并配合轻量级颗粒度先验模块，实现实时设备端交互式分割，达到具有竞争力的准确度。

ABSTRACT

This paper presents EdgeSAM, an accelerated variant of the Segment Anything Model (SAM), optimized for efficient execution on edge devices with minimal compromise in performance. Our approach involves distilling the original ViT-based SAM image encoder into a purely CNN-based architecture, better suited for edge devices. We carefully benchmark various distillation strategies and demonstrate that task-agnostic encoder distillation fails to capture the full knowledge embodied in SAM. To overcome this bottleneck, we include both the prompt encoder and mask decoder in the distillation process, with box and point prompts in the loop, so that the distilled model can accurately capture the intricate dynamics between user input and mask generation. To mitigate dataset bias issues stemming from point prompt distillation, we incorporate a lightweight module within the encoder. As a result, EdgeSAM achieves a 37-fold speed increase compared to the original SAM, and it also outperforms MobileSAM/EfficientSAM, being over 7 times as fast when deployed on edge devices while enhancing the mIoUs on COCO and LVIS by 2.3/1.5 and 3.1/1.6, respectively. It is also the first SAM variant that can run at over 30 FPS on an iPhone 14. Code and demo are available at https://www.mmlab-ntu.com/project/edgesam.

研究动机与目标

• 通过在边缘设备上实现类似 SAM 的功能，激发设备端交互式分割，且不产生巨大的计算成本。
• 研究编码器蒸馏策略并展示任务无关方法的局限性。
• 提出具备提示感知的知识蒸馏，采用动态的提示环路机制。
• 引入轻量级的颗粒度先验模块，以处理数据集特定的提示颗粒度。
• 评估零-shot 转移和在标准分割基准上的实时性能。

提出的方法

• 将 SAM 的 ViT 基于图像编码器蒸馏为纯 CNN 架构，以便边缘部署。
• 应用仅编码器的知识蒸馏，并展示在没有任务特定指导时的局限性。
• 引入动态提示环路蒸馏，在学生出错的区域迭代采样提示以引导掩码解码。
• 保留 SAM 的掩码解码器，并用聚焦于解码器输出的蒸馏目标进行训练。
• 增添轻量级颗粒度先验模块（RPN 与 FPN）以在推理阶段整合数据集特定的颗粒度先验。
• 进行阶段性训练：仅编码器 KD、提示环路 KD，以及可选地使用轻量级 RPN 进行 grounding。

实验结果

研究问题

• RQ1 是否可以通过将 SAM 的组件蒸馏成 CNN 基础编码器，在不显著损失分割质量的情况下将其部署到边缘设备？
• RQ2 在蒸馏过程中使用任务相关的提示引导是否优于仅编码器蒸馏，适用于像 SAM 这样的密集预测任务？
• RQ3 提示类型和颗粒度先验如何影响零-shot 迁移能力以及设备端性能？
• RQ4 哪种骨干网络/骨干网络调整最能在边缘硬件上平衡吞吐量和准确性？
• RQ5 使用基于颗粒度先验的轻量级 RPN 是否能提升对模糊提示的性能？

主要发现

模型	FPS (2080 Ti)	FPS (iPhone 14)	MParam.	GFLOPs
SAM	4.3	N/A	641.1	2734.8
MobileSAM	111.7	5.2	9.8	38.2
EdgeSAM	185.9	72.3	9.6	22.1
EdgeSAM*	135.8	57.8	9.8	22.2

• EdgeSAM 在 iPhone 14 上实现超过 30 FPS，是首个在边缘设备上实现实时性能的 SAM 变体。
• 仅编码器蒸馏本身相对于原始 SAM 存在性能差距，推动需要提示感知蒸馏。
• 提示环路蒸馏在掩码准确性上优于仅编码器 KD，尤其在增加的细化提示下。
• 纯 CNN 骨干（RepViT-M1 with FPN）在边缘部署中提供最佳吞吐-准确性平衡。
• EdgeSAM 在边缘设备上相对于 SAM 提速多达 40x，相对于 MobileSAM 提速多达 14x，在 COCO 和 LVIS 上获得具有竞争力或更好的 mIoU/IoU 增益。
• 以 GT 框作为提示时，EdgeSAM 弥合了与 SAM 的差距，在若干设置中甚至超过 MobileSAM；颗粒度先验的 RPN 进一步提升了 COCO 的中心点性能。