Mambavision L3 512 21K

datasets:

• ILSVRC/imagenet-21k license: other license_name: nvclv1 license_link: LICENSE pipeline_tag: image-classification library_name: transformers

MambaVision: ハイブリッドMamba-Transformer視覚バックボーン.

プロジェクトページ

モデル概要

私たちは、MambaとTransformerの強みを活かした初のハイブリッドコンピュータビジョンモデルを開発しました。具体的には、視覚特徴の効率的なモデリング能力を強化するためにMambaの定式化を再設計しました。さらに、Vision Transformer（ViT）とMambaを統合する可能性について包括的なアブレーション研究を実施しました。結果として、最終層にいくつかのセルフアテンションブロックを備えたMambaアーキテクチャが、長距離空間依存関係を捉えるモデリング能力を大幅に向上させることが示されました。これらの発見に基づき、様々な設計基準を満たす階層型アーキテクチャのMambaVisionモデルファミリーを導入します。

モデル性能

MambaVision-L3-512-21KはImageNet-21Kデータセットで事前学習され、512 x 512解像度でImageNet-1Kでファインチューニングされました。

さらに、MambaVisionモデルはTop-1精度とスループットの面で新しいSOTAパレートフロントを達成し、強力な性能を発揮しています。

モデルの使用方法

MambaVisionの要件をインストールするために、以下を実行することを強く推奨します：

pip install mambavision

各モデルについて、画像分類と特徴抽出の2つのバリアントを提供しており、1行のコードでインポートできます。

画像分類

以下の例では、MambaVisionを画像分類に使用する方法を示します。

COCOデータセットのvalセットから次の画像を入力として与えます：

以下のスニペットを画像分類に使用できます：

from transformers import AutoModelForImageClassification
from PIL import Image
from timm.data.transforms_factory import create_transform
import requests

model = AutoModelForImageClassification.from_pretrained("nvidia/MambaVision-L3-512-21K", trust_remote_code=True)

# 推論用にevalモードに設定
model.cuda().eval()

# モデルのための画像を準備
url = 'http://images.cocodataset.org/val2017/000000020247.jpg'
image = Image.open(requests.get(url, stream=True).raw)
input_resolution = (3, 512, 512)  # MambaVisionは任意の入力解像度をサポート

transform = create_transform(input_size=input_resolution,
                             is_training=False,
                             mean=model.config.mean,
                             std=model.config.std,
                             crop_mode=model.config.crop_mode,
                             crop_pct=model.config.crop_pct)

inputs = transform(image).unsqueeze(0).cuda()
# モデル推論
outputs = model(inputs)
logits = outputs['logits'] 
predicted_class_idx = logits.argmax(-1).item()
print("予測クラス:", model.config.id2label[predicted_class_idx])

予測ラベルはヒグマ、ブラウン・ベア、Ursus arctosです。

特徴抽出

MambaVisionは汎用的な特徴抽出器としても使用できます。

具体的には、モデルの各ステージ（4ステージ）の出力と、最終的な平均プール特徴量を平坦化して抽出できます。

以下のスニペットを特徴抽出に使用できます：

from transformers import AutoModel
from PIL import Image
from timm.data.transforms_factory import create_transform
import requests

model = AutoModel.from_pretrained("nvidia/MambaVision-L3-512-21K", trust_remote_code=True)

# 推論用にevalモードに設定
model.cuda().eval()

# モデルのための画像を準備
url = 'http://images.cocodataset.org/val2017/000000020247.jpg'
image = Image.open(requests.get(url, stream=True).raw)
input_resolution = (3, 512, 512)  # MambaVisionは任意の入力解像度をサポート

transform = create_transform(input_size=input_resolution,
                             is_training=False,
                             mean=model.config.mean,
                             std=model.config.std,
                             crop_mode=model.config.crop_mode,
                             crop_pct=model.config.crop_pct)
inputs = transform(image).unsqueeze(0).cuda()
# モデル推論
out_avg_pool, features = model(inputs)
print("平均プール特徴量のサイズ:", out_avg_pool.size())  # torch.Size([1, 1568])
print("抽出された特徴量のステージ数:", len(features)) # 4ステージ
print("ステージ1の抽出特徴量サイズ:", features[0].size()) # torch.Size([1, 196, 128, 128])
print("ステージ4の抽出特徴量サイズ:", features[3].size()) # torch.Size([1, 1568, 16, 16])

ライセンス:

NVIDIA Source Code License-NC

結果 + 事前学習済みモデル

ImageNet-21K

ImageNet-1K

インストール

Dockerファイルを提供しています。さらに、最近のPyTorchパッケージがインストールされていると仮定すると、依存関係は以下を実行してインストールできます：

pip install -r requirements.txt