许可证：其他 许可证名称：nvclv1 许可证链接：LICENSE 数据集：

• ILSVRC/imagenet-1k 任务标签：图像特征提取

MambaVision：混合Mamba-Transformer视觉骨干网络

模型概述

我们开发了首个结合Mamba与Transformer优势的计算机视觉混合模型。核心贡献包括重新设计Mamba公式以增强其高效建模视觉特征的能力，并通过全面消融实验验证了视觉Transformer（ViT）与Mamba融合的可行性。结果表明，在Mamba架构最后几层加入自注意力模块能显著提升长程空间依赖建模能力。基于此发现，我们推出了具有分层架构的MambaVision系列模型，满足多样化设计需求。

模型性能

MambaVision在Top-1准确率与吞吐量方面实现了新的SOTA帕累托前沿。

使用方法

建议通过以下命令安装依赖环境：

pip install mambavision

每个模型提供图像分类和特征提取两种变体，一行代码即可导入。

图像分类示例

以下演示如何使用MambaVision对COCO数据集验证集图片进行分类：

分类代码片段：

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

model = AutoModelForImageClassification.from_pretrained("nvidia/MambaVision-T2-1K", 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, 224, 224)  # 支持任意输入分辨率

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])

输出结果为：棕熊（brown bear, bruin, Ursus arctos）

特征提取示例

MambaVision可作为通用特征提取器，支持获取模型四个阶段的输出及最终平均池化特征：

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

model = AutoModel.from_pretrained("nvidia/MambaVision-T2-1K", 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, 224, 224)

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, 640])
print("特征阶段数:", len(features)) # 4阶段
print("第一阶段特征尺寸:", features[0].size()) # torch.Size([1, 80, 56, 56])
print("第四阶段特征尺寸:", features[3].size()) # torch.Size([1, 640, 7, 7])

许可证

NVIDIA受限源代码许可证