[논문 리뷰] RepViT-SAM: Towards Real-Time Segmenting Anything

Ao Wang, Hui Chen외 3명|arXiv (Cornell University)|2023. 12. 10.

Advanced Neural Network Applications인용 수 4

한 줄 요약

RepViT-SAM은 Segment Anything Model(SAM)의 무거운 ViT 이미지 인코더를 모바일 최적화된 비전 트랜스포머인 RepViT-M2.3로 대체한다. 이는 효율적인 컨볼루션 아키텍처와 구조적 재패러미터라이제이션을 사용하여 MobileSAM보다 거의 10배 빠른 추론 속도를 달성하면서도 다양한 벤치마크에서 동일하거나 그 이상의 제로샷 세그멘테이션 성능을 유지한다.

ABSTRACT

Segment Anything Model (SAM) has shown impressive zero-shot transfer performance for various computer vision tasks recently. However, its heavy computation costs remain daunting for practical applications. MobileSAM proposes to replace the heavyweight image encoder in SAM with TinyViT by employing distillation, which results in a significant reduction in computational requirements. However, its deployment on resource-constrained mobile devices still encounters challenges due to the substantial memory and computational overhead caused by self-attention mechanisms. Recently, RepViT achieves the state-of-the-art performance and latency trade-off on mobile devices by incorporating efficient architectural designs of ViTs into CNNs. Here, to achieve real-time segmenting anything on mobile devices, following MobileSAM, we replace the heavyweight image encoder in SAM with RepViT model, ending up with the RepViT-SAM model. Extensive experiments show that RepViT-SAM can enjoy significantly better zero-shot transfer capability than MobileSAM, along with nearly $10 imes$ faster inference speed. The code and models are available at \url{https://github.com/THU-MIG/RepViT}.

연구 동기 및 목표

• 자원 제약이 있는 모바일 디바이스에서 Segment Anything Model(SAM)의 실시간 추론을 가능하게 하기 위해.
• MobileSAM과 같은 경량 SAM 버전에서 자기주의 주의 메커니즘의 높은 메모리 및 계산 오버헤드 문제를 해결하기 위해.
• 다양한 세그멘테이션 작업에 걸쳐 경량 SAM 모델의 제로샷 전이 성능을 향상시키기 위해.
• 대규모 ViT-H 인코더에서 더 작은 효율적인 RepViT 기반 인코더로 지식을 전이하기 위해 지식 정복을 활용하기 위해.
• 강력한 성능과 낮은 지연 시간을 갖춘 엣지 배포를 위한 새로운 모바일 최적화 베이스라인을 수립하기 위해.

제안 방법

• SAM의 원래 ViT-H 이미지 인코더를 RepViT-M2.3로 교체한다. 이는 구조적 재패러미터라이제이션과 디프워즈 분리 컨볼루션을 갖춘 경량이고 컨볼루션 아키텍처이다.
• 분리된 지식 정복을 적용한다: ViT-H 인코더에서 RepViT-M2.3 인코더로 특징 맵의 평균 제곱오차(MSE) 손실을 사용해 직접 정복한다.
• RepViT의 최종 다운샘플링 레이어의 스트라이드를 1로 설정하여 SAM의 프롬프트 기반 마스크 디코더와의 출력 해상도 호환성을 확보한다.
• 정복 과정 중에 ViT-H 인코더에서 이미지 임베딩을 사전 계산하고 캐시하여 중량이 큰 인코더를 다시 실행하지 않고도 훈련을 가속화한다.
• 원래 SAM과 동일한 프롬프트 기반 마스크 디코더를 사용하여 추론 파이프라인의 일관성을 유지한다.
• 제로샷 엣지 검출을 위해 16×16 포인트 프롬프트 그리드를 사용하고, 표준 후처리(예: 소벨 필터링 및 비최대 억제)를 적용한다.

Figure 1 : Mask predictions of SAM, MobileSAM, and RepViT-SAM with point prompts (top) and box prompts (bottom).

실험 결과

연구 질문

• RQ1 RepViT처럼 경량이고 컨볼루션 기반의 비전 트랜스포머가 SAM에서 실시간 추론을 달성하면서도 강력한 제로샷 일반화 성능을 유지할 수 있는가?
• RQ2 RepViT-SAM의 성능은 다양한 제로샷 세그멘테이션 벤치마크에서 MobileSAM과 ViT-B-SAM에 비해 어떻게 비교되는가?
• RQ3 대규모 ViT-H 인코더에서 지식 정복을 통해 고품질 특징을 더 작은 RepViT 기반 인코더로 얼마나 효과적으로 전이할 수 있는가?
• RQ4 RepViT의 순수 컨볼루션 설계가 모바일 플랫폼에서 주로 주의 기반 경량 인코더보다 더 나은 지연 시간과 메모리 효율성을 제공하는가?
• RQ5 더 나은 특징 정복 덕분에 RepViT-SAM이 이상 탐지와 같은 특정 후행 작업에서 원래 ViT-H-SAM을 초월할 수 있는가?

주요 결과

• iPhone 12에서 RepViT-SAM은 48.9ms의 추론 지연 시간을 기록했고, MobileSAM은 메모리 초과(OOM) 문제로 실패하여 이는 모바일 배포 가능성의 증거가 된다.
• MacBook M1 Pro에서 RepViT-SAM은 44.8ms로 작동하며, 이는 MobileSAM(482.2ms)보다 약 10배 빠르고 ViT-H-SAM(6249.5ms)보다는 140배 빠르다.
• 제로샷 엣지 검출에서 RepViT-SAM은 ODS 0.764, OIS 0.786, AP 0.773를 기록하여 MobileSAM(0.756, 0.768, 0.746)을 초월하고 ViT-H-SAM(0.768, 0.786, 0.794)의 OIS 성능과 동등하다.
• COCO에서의 제로샷 인스턴스 세그멘테이션에서 RepViT-SAM은 평균 AP 46.1을 기록하여 MobileSAM(43.9)과 ViT-B-SAM(44.8)을 초월하고 ViT-H-SAM(48.7)에 가까워졌다.
• DAVIS 2017에서의 제로샷 비디오 오브제クト 세그멘테이션에서 RepViT-SAM은 평균 J&F 73.5를 기록하여 MobileSAM(71.1)과 ViT-B-SAM(71.3)을 초월했고, MobileSAM보다 2.4포인트 향상되었다.
• MVTec-AD에서의 제로샷 이상 탐지에서 RepViT-SAM은 F1-픽셀 스코어 37.96을 기록하여 ViT-H-SAM(37.65)을 초월하고 MobileSAM(36.44)보다 뛰어나, 특수 작업에서 뛰어난 특징 전이 능력을 보였다.

Figure 2 : Visualization results of SAM, MobileSAM, and RepViT-SAM for zero-shot edge detection on BSDS500.