🚀 MambaVision：混合Mamba-Transformer視覺骨幹網絡

MambaVision是首個用於計算機視覺的混合模型，它結合了Mamba和Transformer的優勢，重新設計了Mamba公式以高效建模視覺特徵，還對Vision Transformers (ViT)與Mamba的集成進行了全面研究，推出了具有分層架構的模型家族，在圖像分類等任務中表現出色。

🚀 快速開始

安裝依賴

強烈建議通過運行以下命令來安裝MambaVision的依賴項：

pip install mambavision

✨ 主要特性

• 混合架構創新：開發了首個結合Mamba和Transformer優勢的計算機視覺混合模型。
• Mamba公式改進：重新設計Mamba公式，增強其對視覺特徵的高效建模能力。
• 集成可行性研究：對Vision Transformers (ViT)與Mamba的集成可行性進行了全面的消融研究。
• 長程依賴建模：在Mamba架構的最後幾層配備多個自注意力塊，顯著提高了捕捉長程空間依賴的建模能力。
• 分層架構模型家族：推出具有分層架構的MambaVision模型家族，以滿足各種設計標準。
• 卓越性能表現：在Top-1準確率和吞吐量方面達到了新的SOTA Pareto前沿。

📦 安裝指南

通過以下命令安裝MambaVision：

pip install mambavision

💻 使用示例

基礎用法

圖像分類

在以下示例中，我們展示瞭如何使用MambaVision進行圖像分類。以COCO數據集驗證集中的圖像作為輸入：

使用以下代碼片段進行圖像分類： ```python from transformers import AutoModelForImageClassification from PIL import Image from timm.data.transforms_factory import create_transform import requests

model = AutoModelForImageClassification.from_pretrained("nvidia/MambaVision-L2-1K", trust_remote_code=True)

eval mode for inference

model.cuda().eval()

prepare image for the model

url = 'http://images.cocodataset.org/val2017/000000020247.jpg' image = Image.open(requests.get(url, stream=True).raw) input_resolution = (3, 224, 224) # MambaVision supports any input resolutions

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

inputs = transform(image).unsqueeze(0).cuda()

model inference

outputs = model(inputs) logits = outputs['logits'] predicted_class_idx = logits.argmax(-1).item() print("Predicted class:", model.config.id2label[predicted_class_idx])

預測標籤為 ```brown bear, bruin, Ursus arctos.```

#### 特徵提取
MambaVision也可以用作通用特徵提取器。可以提取模型每個階段（共4個階段）的輸出以及最終的平均池化特徵。
```python
from transformers import AutoModel
from PIL import Image
from timm.data.transforms_factory import create_transform
import requests

model = AutoModel.from_pretrained("nvidia/MambaVision-L2-1K", trust_remote_code=True)

# eval mode for inference
model.cuda().eval()

# prepare image for the model
url = 'http://images.cocodataset.org/val2017/000000020247.jpg'
image = Image.open(requests.get(url, stream=True).raw)
input_resolution = (3, 224, 224)  # MambaVision supports any input resolutions

transform = create_transform(input_size=input_resolution,
                             is_training=False,
                             mean=model.config.mean,
                             std=model.config.std,
                             crop_mode=model.config.crop_pct,
                             crop_pct=model.config.crop_pct)
inputs = transform(image).unsqueeze(0).cuda()
# model inference
out_avg_pool, features = model(inputs)
print("Size of the averaged pool features:", out_avg_pool.size())  # torch.Size([1, 640])
print("Number of stages in extracted features:", len(features)) # 4 stages
print("Size of extracted features in stage 1:", features[0].size()) # torch.Size([1, 80, 56, 56])
print("Size of extracted features in stage 4:", features[3].size()) # torch.Size([1, 640, 7, 7])

📚 詳細文檔

模型概述

我們開發了首個用於計算機視覺的混合模型，該模型利用了Mamba和Transformer的優勢。具體而言，我們的核心貢獻包括重新設計Mamba公式，以增強其對視覺特徵的高效建模能力。此外，我們對Vision Transformers (ViT)與Mamba的集成可行性進行了全面的消融研究。結果表明，在Mamba架構的最後幾層配備多個自注意力塊，極大地提高了捕捉長程空間依賴的建模能力。基於這些發現，我們推出了具有分層架構的MambaVision模型家族，以滿足各種設計標準。

模型性能

MambaVision表現出色，在Top-1準確率和吞吐量方面達到了新的SOTA Pareto前沿。

📄 許可證

本項目遵循NVIDIA Source Code License-NC許可協議。

相關鏈接

• 論文鏈接
• 代碼倉庫