[논문 리뷰] Faster Segment Anything: Towards Lightweight SAM for Mobile Applications

Chaoning Zhang, Dongshen Han외 5명|arXiv (Cornell University)|2023. 06. 25.

Advanced Neural Network Applications인용 수 141

한 줄 요약

이 논문은 지식 증류를 통해 무거운 이미지 인코더를 대체하는 경량의 모바일 친화적 변형인 MobileSAM을 제시합니다. 이는 원래의 SAM과 FastSAM에 비해 훨씬 작고 더 빠른 추론 속도로 동등한 세분화 성능을 달성합니다.

ABSTRACT

Segment Anything Model (SAM) has attracted significant attention due to its impressive zero-shot transfer performance and high versatility for numerous vision applications (like image editing with fine-grained control). Many of such applications need to be run on resource-constraint edge devices, like mobile phones. In this work, we aim to make SAM mobile-friendly by replacing the heavyweight image encoder with a lightweight one. A naive way to train such a new SAM as in the original SAM paper leads to unsatisfactory performance, especially when limited training sources are available. We find that this is mainly caused by the coupled optimization of the image encoder and mask decoder, motivated by which we propose decoupled distillation. Concretely, we distill the knowledge from the heavy image encoder (ViT-H in the original SAM) to a lightweight image encoder, which can be automatically compatible with the mask decoder in the original SAM. The training can be completed on a single GPU within less than one day, and the resulting lightweight SAM is termed MobileSAM which is more than 60 times smaller yet performs on par with the original SAM. For inference speed, With a single GPU, MobileSAM runs around 10ms per image: 8ms on the image encoder and 4ms on the mask decoder. With superior performance, our MobileSAM is around 5 times faster than the concurrent FastSAM and 7 times smaller, making it more suitable for mobile applications. Moreover, we show that MobileSAM can run relatively smoothly on CPU. The code for our project is provided at \href{https://github.com/ChaoningZhang/MobileSAM}{ extcolor{red}{MobileSAM}}), with a demo showing that MobileSAM can run relatively smoothly on CPU.

연구 동기 및 목표

• 자원 제약이 있는 기기에서 SAM의 모바일 배치를 촉진한다.
• 무거운 이미지 인코더를 교체하여 모델 크기를 줄이고 추론 속도를 향상시킨다.
• 대규모 재학습 없이 원래 SAM의 마스크 디코더와의 호환성을 유지한다.
• 지식 증류 기반 학습이 경량이면서도 정확한 세분화 모델을 만들 수 있음을 보여준다.

제안 방법

• 교사 ViT-H로부터의 지식 증류를 통해 SAM의 ViT-H 이미지 인코더를 경량 인코더로 교체한다.
• 원래의 마스크 디코더를 고정하거나 가볍게 미세조정한 채로 두는 한편, 소형 학생 인코더로 지식을 전달하기 위해 분리된 지식 증류를 적용한다.
• 관련 연구에서 사용된 결합된 focal 및 dice 손실 대신 이미지 임베딩 정렬에 대해 MSE 손실을 사용한다.
• 선택적으로 마스크 디코더를 미세 조정하되, 분리된 지식 증류가 이미 학생 인코더와 디코더 간의 우수한 정합성을 보여줌을 입증한다.
• mIoU 및 추론 속도 메트릭을 사용하여 MobileSAM을 원래의 SAM 및 FastSAM과 비교 평가한다.

실험 결과

연구 질문

• RQ1 무거운 SAM 인코더로부터 증류된 경량 이미지 인코더가 원래 SAM과 유사한 세분화 품질을 달성할 수 있는가?
• RQ2 분리된 지식 증류가 경량 SAM 학습에서 결합 증류 또는 반결합 증류보다 우수한가?
• RQ3 Segment-anything 작업에서 정확도(mIoU)와 효율성(매개변수 수, 속도) 측면에서 MobileSAM은 FastSAM과 어떻게 비교되는가?
• RQ4 MobileSAM이 온-디바이스 애플리케이션을 위해 CPU에서 효율적으로 실행될 수 있는가?

주요 결과

• MobileSAM은 인코더 매개변수를 약 100배, 전체 매개변수를 약 60배 감소시키면서 원래의 SAM과 유사한 성능을 달성합니다.
• 단일 GPU에서 MobileSAM은 약 10 ms(인코더 8 ms, 디코더 4 ms) 동안 이미지를 처리합니다.
• MobileSAM은 FastSAM보다 대략 5배 빠르고 7배 더 작으며, Segment Anything 설정에서 우수한 성능을 제공합니다.
• 분리된 지식 증류가 예비 실험에서 결합 증류보다 더 나은 mIoU(0.75)를 나타냅니다.
• MobileSAM은 CPU에서 비교적 부드럽게 실행되어 모바일 기기 배치를 가능하게 합니다.