MambaVision-T-1K

MambaVision-T-1K项目介绍

项目概述

MambaVision-T-1K项目是一个创新的计算机视觉模型，结合了Mamba和Transformer的优势，首次提出了混合模型的概念。这个项目的核心贡献在于重新设计了Mamba框架，以提升其视觉特征建模的效率。此外，项目团队对将视觉Transformer（ViT）与Mamba结合的可行性进行了全面的消融研究。研究结果表明，在Mamba架构的最后几层引入若干自注意力模块，能够大幅提高模型捕捉长距离空间依赖的能力。基于此，项目组推出了一系列分层结构的MambaVision模型，以满足不同的设计需求。

模型表现

MambaVision在Top-1准确率和吞吐量方面实现了新的最优性能前端，展示了其强大的性能。

模型使用

推荐用户通过以下命令来安装MambaVision的必要依赖：

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

预测标签为“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-T-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)  # 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, 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])

许可协议

模型使用NVIDIA源代码许可协议-NC，具体信息请访问NVIDIA Source Code License-NC。