EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM
Chong Zhou Xiangtai Li Chen Change Loy Bo Dai
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. EdgeSAM achieves a 40-fold speed
increase compared to the original SAM, and it also outperforms MobileSAM, being
14 times as fast when deployed on edge devices while enhancing the mIoUs on
COCO and LVIS by 2.3 and 3.2 respectively. It is also the first SAM variant
that can run at over 30 FPS on an iPhone 14. Code and models are available at
https://github.com/chongzhou96/EdgeSAM.