[EdgeSAM] 논문 리뷰

 

SAM의 경량화 모델인 EdgeSAM에 관한 논문을 리뷰해 보았다.

 

Introduction

 

이미지 내 객체 영역을 추출하는 segmentation 모델인 SAM(segment anything model)의 한계점 중 하나는 segment를 수행하는 Encoder가 무겁다는 점이다. EdgeSAM은 SAM에 비해 파라미터 개수를 0.6%로 줄이고 데이터셋의 1% 만을 사용한 모델로 encoder의 속도가 빠르다. 이는 서비스에 적용하기에 더욱 적합한 모델임을 뜻하며, 기존 SAM에 비해 40배 빠르며 아이폰 14에서 실시간 제공이 가능하다고 한다.

---

 

Method

 

prompt-in-the-loop distillation

EdgeSAM의 학습 방식인 prompt-in-the-loop distillation에 대해 아래에 그림과 함께 설명하였다.

 

[그림 1] prompt-in-the-loop distillation

 

 

[그림 1]의 prompt-in-the-loop distillation은 딥러닝 모델의 압축 또는 경량화를 위한 기술로 Decoder는 고정하고 SAM과 EdgeSAM의 encoder output loss를 줄이는 방식으로 학습하는 방식이다. 이는 SAM과 EdgeSAM의 객체 추출 결과가 유사해지도록 학습하는 방식이다. EdgeSAM은 이 학습 방식을 사용하여 Encoder를 기존 SAM에 비해 경량화할 수 있다.

 

[그림 2] 학습 샘플 추출 방법

 

EdgeSAM Encoder는 SAM과 EdgeSAM의 segment 결과에서 FN (False negative,  EdgeSAM은 아니라고 판단했지만 SAM은 맞다고 판단), FP (False positive, EdgeSAM은 맞다고 판단했지만 SAM은 아니라고 판단) point를 랜덤으로 추출하여 반복적으로 학습한다. [그림 2]에서 Teacher Mask가 SAM, Student Mask가 EdgeSAM을 의미한다.

---

 

Experiments

 

데이터셋은 SA-1B, COCO, LVIS를 사용하였다. distillation을 위해 SA-1B 데이터셋의 1%만 사용하여 SA-1K 데이터셋을 만들었으며, training에 쓰이지 않은 1000개를 랜덤으로 추출하여 테스트셋으로 사용하였다. 데스크탑과 iphone 14 환경 모두에서 테스트 해보았을 때, EdgeSAM이 SAM과 MobileSAM보다 처리 속도가 월등히 빨랐다. 

성능 평가를 위해서는 3가지 방법을 사용하였다. 첫 번째로 ground-truth box를 초기 prompt로 하여 반복적으로 추가적인 point를 찾아나가는 방법이다. 두 번째로 ground-truth box의 중앙 point를 초기 prompt로 하여 반복적으로 추가적인 point를 찾아나가는 방법이다. 마지막 세 번째로 외부 객체 detector(논문에서는 Detic을 사용함)로 초기 prompt를 설정하여 추가적인 point를 찾아나가는 방법이다.

첫 번째, 두 번째 방법의 성능 확인을 위해서는 mIoU(모든 클래스에 대한 객체 겹침 정도의 평균값), 세 번째 방법의 성능 확인을 위해서는 mAP(탐지 결과의 정밀도의 평균값)와 boundary IoU( ground-truth box와 실제 라벨 간의 겹침 정도)를 사용하였다. 전반적으로 SAM보다 정확도는 조금 떨어졌지만, FPS(초당 처리 프레임 수)가 SAM에 비해 매우 빨랐다.

 

논문 출처: Chong Zhou et al.,  EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM, 2023, https://arxiv.org/pdf/2312.06660.pdf