[論文レビュー] RepViT-SAM: Towards Real-Time Segmenting Anything

Ao Wang, Hui Chen+3 more|arXiv (Cornell University)|Dec 10, 2023

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 のプロンプト駆動マスクデコーダーとの出力解像度の整合性を保証する。
• 蒸留中に、重いエンコーダーを再実行せずに、画像埋め込みを事前に計算してキャッシュする。
• 元の SAM と同一のプロンプト駆動マスクデコーダーを採用し、推論パイプラインの一貫性を維持する。
• ゼロショットエッジ検出には 16×16 の前景ポイントプロンプトグリッドを用い、Sobelフィルタリングと非最大抑制を含む標準的な後処理を適用する。

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）を上回り、OIS では ViT-H-SAM（0.768、0.786、0.794）と同等の性能を示した。
• COCO におけるゼロショットインスタンスセグメンテーションでは、RepViT-SAM が 46.1 の平均 AP を達成し、MobileSAM（43.9）と ViT-B-SAM（44.8）を上回り、ViT-H-SAM（48.7）に近づいた。
• DAVIS 2017 におけるゼロショット動画オブジェクトセグメンテーションでは、RepViT-SAM が 73.5 の平均 J&F を達成し、MobileSAM（71.1）と ViT-B-SAM（71.3）を上回り、MobileSAM よりも 2.4 ポints の向上を示した。
• MVTec-AD におけるゼロショット異常検出では、RepViT-SAM が 37.96 の F1ピクセルスコアを達成し、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.