🚀 MambaVision:混合Mamba-Transformer視覺骨幹網絡
MambaVision是首個用於計算機視覺的混合模型,結合了Mamba和Transformer的優勢,能高效建模視覺特徵,在圖像分類任務中表現出色。
🚀 快速開始
安裝
強烈建議通過運行以下命令來安裝MambaVision所需的依賴:
pip install mambavision
✨ 主要特性
- 開發了首個結合Mamba和Transformer優勢的計算機視覺混合模型。
- 重新設計Mamba公式,增強其對視覺特徵的高效建模能力。
- 對Vision Transformers (ViT) 與Mamba集成的可行性進行了全面的消融研究。
- 在Mamba架構的最後幾層配備多個自注意力塊,顯著提高了捕捉長距離空間依賴的建模能力。
- 引入了具有分層架構的MambaVision模型家族,以滿足各種設計標準。
📚 詳細文檔
模型概述
我們開發了首個用於計算機視覺的混合模型,充分利用了Mamba和Transformer的優勢。具體而言,我們的核心貢獻包括重新設計Mamba公式,以增強其對視覺特徵進行高效建模的能力。此外,我們對將視覺Transformer(ViT)與Mamba集成的可行性進行了全面的消融研究。結果表明,在Mamba架構的最後幾層配備多個自注意力塊,極大地提高了捕捉長距離空間依賴的建模能力。基於這些發現,我們引入了具有分層架構的MambaVision模型家族,以滿足各種設計標準。
模型性能
MambaVision-L3-256-21K在ImageNet-21K數據集上進行預訓練,並在ImageNet-1K上進行微調。預訓練和微調均在256 x 256分辨率下進行。
名稱 |
準確率@1(%) |
準確率@5(%) |
參數數量(M) |
浮點運算次數(G) |
分辨率 |
MambaVision-L3-256-21K |
87.3 |
98.3 |
739.6 |
122.3 |
256x256 |
此外,MambaVision模型在Top-1準確率和吞吐量方面達到了新的最優帕累託前沿,表現出色。
💻 使用示例
基礎用法 - 圖像分類
以下示例展示瞭如何使用MambaVision進行圖像分類:
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-256-21K", trust_remote_code=True)
model.cuda().eval()
url = 'http://images.cocodataset.org/val2017/000000020247.jpg'
image = Image.open(requests.get(url, stream=True).raw)
input_resolution = (3, 256, 256)
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])
預測標籤為 棕熊,bruin,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-256-21K", trust_remote_code=True)
model.cuda().eval()
url = 'http://images.cocodataset.org/val2017/000000020247.jpg'
image = Image.open(requests.get(url, stream=True).raw)
input_resolution = (3, 256, 256)
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())
print("提取特徵的階段數:", len(features))
print("第1階段提取特徵的大小:", features[0].size())
print("第4階段提取特徵的大小:", features[3].size())
📄 許可證
本項目遵循 NVIDIA源代碼許可協議 - 非商業用途。
相關鏈接