🚀 MambaVision:混合Mamba - 变压器视觉骨干网络
MambaVision是首个用于计算机视觉的混合模型,它结合了Mamba和Transformer的优势,重新设计了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数据集验证集中的图像作为输入:
使用以下代码片段进行图像分类:
from transformers import AutoModelForImageClassification
from PIL import Image
from timm.data.transforms_factory import create_transform
import requests
model = AutoModelForImageClassification.from_pretrained("nvidia/MambaVision-S-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("Predicted class:", model.config.id2label[predicted_class_idx])
预测标签为 brown bear, 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-S-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("Size of the averaged pool features:", out_avg_pool.size())
print("Number of stages in extracted features:", len(features))
print("Size of extracted features in stage 1:", features[0].size())
print("Size of extracted features in stage 4:", features[3].size())
📚 详细文档
模型概述
我们开发了首个用于计算机视觉的混合模型,该模型利用了Mamba和Transformer的优势。具体而言,我们的核心贡献包括重新设计Mamba公式,以增强其对视觉特征的高效建模能力。此外,我们对Vision Transformers (ViT) 与Mamba集成的可行性进行了全面的消融研究。结果表明,在Mamba架构的最后几层配备多个自注意力块,极大地提高了捕捉长程空间依赖的建模能力。基于这些发现,我们引入了具有分层架构的MambaVision模型家族,以满足各种设计标准。
模型性能
MambaVision表现出色,在Top - 1准确率和吞吐量方面达到了新的SOTA Pareto前沿。
📄 许可证
MambaVision: A Hybrid Mamba - Transformer Vision Backbone遵循NVIDIA Source Code License - NC许可协议。