Mambavision L2 512 21K

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

MambaVision是首個用於計算機視覺的混合模型，結合了Mamba和Transformer的優勢，重新設計Mamba公式以高效建模視覺特徵，還研究了與ViT集成的可行性，推出分層架構的模型家族，滿足不同設計需求。

🚀 快速開始

安裝

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

pip install mambavision

使用示例

基礎用法

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-L2-512-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, 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("Predicted class:", model.config.id2label[predicted_class_idx])

# 特徵提取示例
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-512-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, 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("Size of the averaged pool features:", out_avg_pool.size())  # torch.Size([1, 1568])
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, 196, 128, 128])
print("Size of extracted features in stage 4:", features[3].size()) # torch.Size([1, 1568, 16, 16])

✨ 主要特性

• 開發了首個結合Mamba和Transformer優勢的計算機視覺混合模型。
• 重新設計Mamba公式，增強其對視覺特徵的高效建模能力。
• 對Vision Transformers (ViT) 與Mamba集成的可行性進行了全面的消融研究。
• 提出了具有分層架構的MambaVision模型家族，以滿足各種設計標準。

📦 安裝指南

通過以下命令安裝MambaVision：

pip install mambavision

📚 詳細文檔

模型概述

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

模型性能

MambaVision-L2-512-21K在ImageNet-21K數據集上進行預訓練，並在512 x 512分辨率的ImageNet-1K上進行微調。

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

模型使用

MambaVision可用於圖像分類和特徵提取，具體使用方法見上文的使用示例。

許可證

NVIDIA源代碼許可協議 - 非商業用途

🔧 技術細節

• 提出了一種混合模型，結合了Mamba和Transformer的優勢，用於計算機視覺任務。
• 重新設計了Mamba公式，以提高其對視覺特徵的建模能力。
• 研究了Vision Transformers (ViT) 與Mamba集成的可行性，並通過實驗證明了在Mamba架構的最後幾層添加自注意力塊可以提高其捕捉長距離空間依賴關係的能力。
• 推出了具有分層架構的MambaVision模型家族，以滿足不同的設計需求。

📄 許可證

本項目採用 NVIDIA源代碼許可協議 - 非商業用途。

信息表格