EfficientViT-SAM: 精度を変えずにその場で離陸!

著者らは、高速化された SAM モデルファミリである EfficientViT-SAM を提案しました。著者らは、SAM の軽量なヒントエンコーダーとマスクデコーダーを保持しながら、重いイメージエンコーダーを EfficientViT に置き換えました。トレーニングでは、まずSAM-ViT-H画像エンコーダーからEfficientViTへの知識蒸留が実行されます。続いて、SA-1B データセットに対してエンドツーエンドのトレーニングが実行されます。 EfficientViT の効率性と容量のおかげで、EfficientViT-SAM はパフォーマンスを犠牲にすることなく、A100 GPU で 48.9% の TensorRT 高速化を実現します。

コードと事前トレーニング: https://github.com/mit-han-lab/efficientvit

1 はじめに

Segment Anything Model (SAM) は、1,100 万枚の画像と 10 億個のマスクを含む高品質のデータセットで事前トレーニングされた一連の画像セグメンテーションモデルです。 SAM は驚異的なゼロショット画像セグメンテーション性能を提供し、拡張現実/仮想現実、データ注釈、インタラクティブな画像編集など、多くのアプリケーションで使用されています。

SAM は強力なパフォーマンスを発揮しますが、計算負荷が非常に高いため、時間に敏感なシナリオでは適用が制限されます。特に、SAM の主な計算ボトルネックはその画像エンコーダーにあり、推論時に画像ごとに 2973 GMAC が必要です。

SAM を高速化するために、SAM の画像エンコーダーを軽量モデルに置き換える試みが数多く行われてきました。例えば、MobileSAM[22]はSAMのViT-Hモデルの知識を小さな視覚的トランスフォーマーに凝縮しています。 EdgeSAM は、ViT-H を模倣するために純粋な CNN ベースのモデルをトレーニングし、ヒントエンコーダーとマスクデコーダーを含む詳細な蒸留戦略を採用します。 EfficientSAM は、MAE 事前トレーニング方式を利用してパフォーマンスを向上させます。

これらの方法は計算コストを削減できますが、いずれもパフォーマンスが大幅に低下します (図 1)。本稿では、SAM の画像エンコーダーを EfficientViT に置き換えることでこの制限に対処するEfficientViT-SAMを紹介します。同時に、著者らは SAM の軽量ヒントエンコーダーとマスクデコーダーのアーキテクチャを保持しました。著者のトレーニングプロセスは 2 つの段階で構成されています。まず、SAM の画像エンコーダーを教師として使用して、EfficientViT-SAM の画像エンコーダーをトレーニングします。次に、著者らは SA-1B データセット全体を使用して、EfficientViT-SAM をエンドツーエンドでトレーニングします。

著者らは、一連のゼロショットベンチマークで EfficientViT-SAM のパフォーマンスを総合的に評価します。 EfficientViT-SAM は、パフォーマンスと効率の両方において、これまでのすべての SAM モデルを大幅に上回ります。特に、COCO データセットでは、EfficientViT-SAM は、mAP を低下させることなく、SAM-ViT-H と比較して A100 GPU で 48.9 倍のスループット向上を実現します。

2 関連研究

セグメント化モデル

SAM は、画像セグメンテーションにおける優れたパフォーマンスと一般化能力を実証しており、コンピュータービジョンの分野では、この分野におけるマイルストーンとして広く認識されています。 SAM は、画像セグメンテーションをヒント可能なタスクとして定義し、セグメンテーションのヒントが与えられた場合に有効なセグメンテーションマスクを生成することを目指しています。この目標を達成するために、SAM は画像エンコーダーとヒントエンコーダーを使用して画像を処理してヒントを提供します。 2 つのエンコーダーの出力はマスクデコーダーに送られ、最終的なマスク予測が生成されます。

SAM は、1,100 万枚を超える画像と 10 億を超える高品質マスクを含む大規模なセグメンテーションデータセットでトレーニングされており、ゼロショットのオープンワールドセグメンテーションにおいて強力な機能を発揮します。 SAM は、画像の修復、オブジェクトの追跡、3D 生成など、さまざまな下流アプリケーションで高い適応性を発揮します。ただし、SAM の画像エンコーダーコンポーネントは計算コストが非常に高く、レイテンシが高くなるため、時間に敏感なシナリオでは実用性が制限されます。最近の研究では、SAM の計算上の限界に対処することを目的として、SAM の効率を改善することに重点が置かれています。

効率的なディープラーニングコンピューティング

ディープニューラルネットワークの効率を向上させることは、エッジおよびクラウドコンピューティングプラットフォームでの実際のアプリケーションにとって非常に重要です。私たちの研究は、効果的なモデルアーキテクチャの設計に関連しており、非効率的なモデルアーキテクチャを効率的なモデルアーキテクチャに置き換えることで、パフォーマンスと効率のトレードオフを改善することを目指しています。著者らの研究は、事前にトレーニングされた教師モデルを利用して生徒モデルのトレーニングをガイドする知識蒸留にも関連しています。さらに、著者らは、EfficientViT-SAM を他の並列技術と組み合わせて、プルーニング、量子化、ハードウェア対応のニューラルアーキテクチャ検索などの効率をさらに向上させることができます。

3 方法

著者らは、EfficientViT を使用して SAM を高速化する EfficientViT-SAM を提案しました。特に、EfficientViT-SAM は、イメージエンコーダを EfficientViT に置き換えながら、SAM のヒントエンコーダとマスクデコーダのアーキテクチャを保持します。著者らは、速度とパフォーマンスのバランスを実現する EfficientViT-SAM-L と EfficientViT-SAM-XL という 2 つのモデルファミリを設計しました。次に著者らは SA-1B データセットを使用して、EfficientViT-SAM をエンドツーエンドでトレーニングしました。

効率的なViT

EfficientViT は、効率的な高解像度の高密度予測を実現するビジュアル Transformer モデルファミリです。その中核となる構成要素は、ハードウェア効率の高い操作を通じてグローバル受容野とマルチスケール学習を実現するマルチスケール線形注意モジュールです。

具体的には、非効率的なソフトマックスアテンションを軽量の ReLU 線形アテンションに置き換えて、グローバル受容フィールドを実現します。行列乗算の結合特性を活用することで、ReLU 線形アテンションは機能性を維持しながら計算の複雑さを 2 次から 1 次へと削減できます。さらに、畳み込みによって ReLU 線形アテンションを強化し、ローカル特徴抽出における制限を軽減します。詳細は原文論文をご覧ください。

効率的なViT-SAM

モデルアーキテクチャ。 EfficientViT-SAM-XL のマクロアーキテクチャを図 2 に示します。本体は5つのステージから構成されます。 EfficientViT と同様に、著者らは初期段階で畳み込みブロックを使用し、最後の 2 段階で efficientViT モジュールを使用します。著者らは、アップサンプリングと追加を通じて最後の 3 つのステージの機能を融合します。融合された機能は、複数の融合された MBConv ブロックで構成される Neck に送られ、その後 SAM Head に送られます。

電車。画像エンコーダーを初期化するには、まず SAM-ViT-H からの画像埋め込み情報を EfficientViT に抽出します。著者は損失関数としてL2損失を使用しています。ヒントエンコーダーとマスクデコーダーについては、SAM-ViT-H の重みを読み込んで初期化します。次に、SA-1B データセットで EfficientViT-SAM をエンドツーエンドでトレーニングします。

エンドツーエンドのトレーニングフェーズでは、ボックスキューとポイントキューを同じ確率でランダムに選択します。ポイントヒントの場合、モデルがさまざまなポイント構成に効果的に対応できるように、実際のマスクから 1 ～ 10 個の前景ポイントをランダムに選択します。ボックスヒントの場合は、グラウンドトゥルース境界ボックスを使用します。 EfficientViT-SAM-L/XL モデルの場合、最長のエッジのサイズを 512/1024 に変更し、それに応じて短いエッジをパディングします。著者は、画像ごとに最大 64 個のランダムにサンプリングされたマスクを選択します。

トレーニングプロセスを監視するために、著者は Focal Loss と Dice Loss の線形結合を使用し、Focal Loss と Dice Loss の比率を 20:1 にします。 SAM で使用される曖昧さ解消法と同様に、著者らは 3 つのマスクを同時に予測し、損失が最も低いマスクのみを逆伝播します。作者は、4 番目の出力トークンを追加することで、単一のマスクの出力もサポートしています。トレーニング中は、2 つの予測モードをランダムに切り替えます。

著者らは、SA-1B データセットを使用して、バッチサイズ 256 で EfficientViT-SAM を 2 エポックにわたってトレーニングしました。 AdamW オプティマイザーは、運動量パラメータを 0.9、確率を 0.999 に設定して使用されました。初期学習率は、EfficientViT-SAM-L/XL の場合はそれぞれ 2e/1e に設定され、コサイン減衰学習率スケジュールを使用して 0 に削減されます。データ増強のために、著者らはランダムな水平反転を適用した。

4 実験

このセクションでは、セクション 4.1 で EfficientViT-SAM の実行時効率の包括的な分析を示します。次に著者らは、トレーニング中には遭遇しなかった COCO および LVIS データセットで EfficientViT-SAM のゼロショット機能を評価しました。著者らは、セクション 4.2 でのシングルショット有効マスク評価とセクション 4.3 での境界ボックスによるインスタンス分割のヒントという 2 つの異なるタスクを実行しました。これらのタスクは、それぞれ EfficientViT-SAM のポイントキュー機能とバウンディングボックスキュー機能の有効性を評価します。さらに、著者らはセクション4.4でSGlnWベンチマークの結果も提供しています。

実行時の効率

著者らは、EfficientViT-SAM のモデルパラメーター、MAC、スループットを SAM およびその他の高速化作業と比較します。結果を表1に示す。著者らは単一の NVIDIA A100 GPU でスループットを測定し、TensorRT 最適化を使用しました。

結果によると、SAM と比較して 17 ～ 69 倍という驚異的な高速化が達成されました。さらに、EfficientViT-SAM は他のアクセラレーション作業よりも多くのパラメータを備えていますが、ハードウェアフレンドリーな演算子を効果的に利用することでスループットが大幅に向上します。

ゼロショットポイントプロンプトセグメンテーション

表 2 では、ポイントキューに基づいてオブジェクトをセグメント化する場合の EfficientViT-SAM のゼロショットパフォーマンスを評価します。著者らは[1]で説明したポイント選択方法を採用した。つまり、初期点はターゲット境界から最も遠い点として選択されます。後続の各ポイントは、真の値と以前の予測値の間の領域として定義されるエラー領域の境界から最も遠いポイントとして選択されます。

著者らは、1/3/5 クリックを使用して COCO および LVIS データセットのパフォーマンスを報告し、mIoU (平均交差和集合) を評価基準としています。私たちの結果は、特に追加のポイントキューを提供する場合に、SAM と比較して優れたパフォーマンスを示しています。

ゼロショットボックスプロンプトセグメンテーション

著者らは、バウンディングボックスを使用したオブジェクトセグメンテーションにおける EfficientViT-SAM のゼロショットパフォーマンスを評価しました。まず、著者らは真の境界ボックスをモデルに入力し、その結果を表 4 に示します。

mIoU (平均交差和集合) は、すべてのオブジェクトについて、また小さいオブジェクト、中くらいのオブジェクト、大きいオブジェクトごとに個別に報告されます。 EfficientViT-SAM は、COCO および LVIS データセットで SAM を大幅に上回ります。次に、著者らはオブジェクト検出器 ViT-Det を採用し、その出力ボックスをモデルのヒントとして使用します。表 5 の結果は、EfficientViT-SAM が SAM よりも優れたパフォーマンスを達成することを示しています。 EfficientViT-SAM の最も軽量なバージョンでさえ、他の高速化作業よりも大幅に優れていることは注目に値します。

さらに、著者らはYOLOv8とGroundingDINOをターゲット検出器として使用し、COCOデータセットでのEfficientViT-SAMのパフォーマンスを評価しました。 YOLOv8 は、実用的なアプリケーションシナリオに適したリアルタイムオブジェクト検出器です。一方、GroundingDINO はテキストの手がかりを使用してオブジェクトを検出できるため、テキストの手がかりに基づいてオブジェクトのセグメンテーションを実行できます。表 6 に示す結果は、EfficientViT-SAM が SAM と比較して優れたパフォーマンスを発揮することを示しています。

ゼロショットのインザワイルドセグメンテーション

インザワイルドセグメンテーションベンチマークには、25 個のゼロショットインザワイルドセグメンテーションデータセットが含まれています。著者らは、ゼロショットセグメンテーションを実行するためのボックスヒントとして、EfficientViT-SAM と Grounding-DINO を組み合わせました。各データセットの全体的なパフォーマンス結果を表 3 に示します。 SAM は mAP 48.7 を達成し、EfficientViT-SAM はさらに高いスコア 48.9 を達成しました。

定性的な結果。

図 3 は、ポイントキュー、ボックスキュー、および SAM モードが提供された場合の EfficientViT-SAM の定性的なセグメンテーション結果を示しています。結果は、EfficientViT-SAM が大きなオブジェクトのセグメント化に優れているだけでなく、小さなオブジェクトも効果的に処理できることを示しています。これらの結果は、EfficientViT-SAM の優れたセグメンテーション能力を浮き彫りにしています。