EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM

Chong Zhou,Xiangtai Li,Chen Change Loy,Bo Dai

2024-07-20

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 taskagnostic 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 <a class="link-external link-https" href="https://www.mmlab-ntu.com/project/edgesam" rel="external noopener nofollow">this https URL</a>.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

### Main Problems Addressed by the Paper

The paper primarily addresses the issues of high computational resource consumption, slow speed, and performance degradation when deploying the Segment Anything Model (SAM) on edge devices. Specifically:

1. **Accelerating the SAM Model**: The paper proposes a new model called EdgeSAM, which aims to transfer the core functionalities of SAM to a more compact architecture suitable for edge device operation through knowledge distillation techniques.

2. **Optimizing Model Structure**: By distilling the image encoder based on Vision Transformer (ViT) in SAM into an architecture purely based on Convolutional Neural Networks (CNN), the execution efficiency on edge devices is improved.

3. **Improving Distillation Strategy**: A method called "Prompt-In-The-Loop Distillation" is introduced to guide the learning process of the student model. This method can more effectively capture the complex dynamics in SAM, especially in handling the relationship between user input and mask generation.

4. **Addressing Dataset Bias**: To solve the potential dataset bias issue during point prompt distillation, a lightweight module is integrated within the encoder.

5. **Enhancing Practical Application Performance**: EdgeSAM significantly improves the running speed on edge devices while maintaining or approaching the original SAM performance. For example, it achieves a speed of over 30 FPS on an iPhone 14 and improves performance across multiple datasets.

Through these methods, EdgeSAM can greatly enhance the real-time interactive segmentation capability of the SAM model on edge devices while ensuring performance, thereby promoting its integration and application in various downstream tasks.