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

MambaVision：一种混合型Mamba-Transformer视觉骨干网络 的官方 PyTorch 实现，详见 arXiv 预印本。

作者：Ali Hatamizadeh 和 Jan Kautz。

如需商务合作，请访问我们的官网并提交表格：NVIDIA Research Licensing

立即体验 MambaVision：

---

MambaVision 在 Top-1 准确率和吞吐量方面取得了新的 SOTA 帕累托前沿，表现出色。

我们引入了一种新颖的混合模块，通过构建无 SSM 的对称路径来增强全局上下文建模能力：

MambaVision 采用层次化架构，同时使用自注意力机制和混合模块：

💥 最新消息 💥

• [2025年6月10日] MambaVision 的海报将亮相 2025 年 CVPR 大会，地点为纳什维尔，展示时间为 2025 年 6月15日星期日 上午10:30 至中午12:30（CDT），展位号为 D 展厅 403 号。

• [2025年6月10日] 语义分割代码及模型已发布 此处！

• [2025年6月7日] 目标检测代码及模型已发布 此处！

• [2025年3月29日] 现在您可以在 Google Colab 中轻松运行 MambaVision。立即尝试：

• [2025年3月29日] 新的 MambaVision pip 包 已发布！

• [2025年3月25日] 更新后的 论文 现已在 arXiv 上公开！

• [2025年3月25日] 仓库中新增了 21K 模型和代码。

• [2025年3月25日] MambaVision 是首个大规模的基于 Mamba 的视觉骨干网络！

• [2025年3月24日] MambaVision-L3-512-21K 达到 88.1% 的 Top-1 准确率。

• [2025年3月24日] 新的 ImageNet-21K 模型已加入 MambaVision Hugging Face 收藏。

• [2025年2月26日] MambaVision 已被 2025 年 CVPR 接受！

• [2024年7月24日] MambaVision Hugging Face 模型正式发布！

• [2024年7月14日] 我们增加了对任意分辨率图像处理的支持。

• [2024年7月12日] 论文 现已在 arXiv 上发布！

• [2024年7月11日] Mambavision pip 包 正式发布！

• [2024年7月10日] 我们发布了 MambaVision 的代码和模型检查点！

快速入门

Google Colab

您可以在 Google Colab 中简单试用 MambaVision 进行图像分类：

Hugging Face（分类 + 特征提取）

预训练的 MambaVision 模型可通过 Hugging Face 库以 几行代码 轻松使用。首先安装依赖：

pip install mambavision

然后即可导入模型：

>>> from transformers import AutoModelForImageClassification

>>> model = AutoModelForImageClassification.from_pretrained("nvidia/MambaVision-T-1K", trust_remote_code=True)

以下是一个端到端的图像分类示例。

输入图片来自 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("预测类别:", model.config.id2label[predicted_class_idx])

预测结果为棕熊、灰熊，学名 Ursus arctos。

您还可以使用 Hugging Face 上的 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("平均池化特征的尺寸:", out_avg_pool.size())  # torch.Size([1, 640])
print("提取特征中的阶段数:", len(features)) # 4个阶段
print("第1阶段提取特征的尺寸:", features[0].size()) # torch.Size([1, 80, 56, 56])
print("第4阶段提取特征的尺寸:", features[3].size()) # torch.Size([1, 640, 7, 7])

目前，我们在 Hugging Face 上提供了 MambaVision-T-1K、MambaVision-T2-1K、MambaVision-S-1K、MambaVision-B-1K、MambaVision-L-1K 和 MambaVision-L2-1K。所有模型也可以在 这里 查看。

分类（pip 包）

我们还可以通过 pip 包以 几行代码 导入预训练的 MambaVision 模型：

pip install mambavision

可以使用默认超参数创建一个预训练的 MambaVision 模型，如下所示：

>>> from mambavision import create_model

# 定义 mamba_vision_T 模型

>>> model = create_model('mamba_vision_T', pretrained=True, model_path="/tmp/mambavision_tiny_1k.pth.tar")

可用的预训练模型包括 mamba_vision_T、mamba_vision_T2、mamba_vision_S、mamba_vision_B、mamba_vision_L 和 mamba_vision_L2。

我们也可以简单地通过传递一张 任意分辨率 的虚拟图像来测试模型。输出是 logits：

>>> import torch

>>> image = torch.rand(1, 3, 512, 224).cuda() # 将图像放在 GPU 上
>>> model = model.cuda() # 将模型放在 GPU 上
>>> output = model(image) # 输出 logit 的尺寸是 [1, 1000]

使用我们 pip 包中的预训练模型，您可以轻松运行验证：

python validate_pip_model.py --model mamba_vision_T --data_dir=$DATA_PATH --batch-size $BS 

结果 + 预训练模型

ImageNet-21K

名称	Acc@1(%)	Acc@5(%)	#参数(M)	FLOPs(G)	分辨率	Hugging Face	下载
MambaVision-B-21K	84.9	97.5	97.7	15.0	224x224	链接	模型
MambaVision-L-21K	86.1	97.9	227.9	34.9	224x224	链接	模型
MambaVision-L2-512-21K	87.3	98.4	241.5	196.3	512x512	链接	模型
MambaVision-L3-256-21K	87.3	98.3	739.6	122.3	256x256	链接	模型
MambaVision-L3-512-21K	88.1	98.6	739.6	489.1	512x512	链接	模型

ImageNet-1K

名称	Top-1准确率(%)	Top-5准确率(%)	吞吐量(张/秒)	分辨率	参数量(M)	FLOPs(G)	Hugging Face链接	下载链接
MambaVision-T	82.3	96.2	6298	224x224	31.8	4.4	链接	模型
MambaVision-T2	82.7	96.3	5990	224x224	35.1	5.1	链接	模型
MambaVision-S	83.3	96.5	4700	224x224	50.1	7.5	链接	模型
MambaVision-B	84.2	96.9	3670	224x224	97.7	15.0	链接	模型
MambaVision-L	85.0	97.1	2190	224x224	227.9	34.9	链接	模型
MambaVision-L2	85.3	97.2	1021	224x224	241.5	37.5	链接	模型

检测结果 + 模型

骨干网络	检测器	学习率调度	边界框mAP	掩码mAP	参数量(M)	FLOPs(G)	配置文件	日志	模型检查点
MambaVision-T-1K	级联Mask R-CNN	3倍	51.1	44.3	86	740	配置文件	日志	模型
MambaVision-S-1K	级联Mask R-CNN	3倍	52.3	45.2	108	828	配置文件	日志	模型
MambaVision-B-1K	级联Mask R-CNN	3倍	52.8	45.7	145	964	配置文件	日志	模型

分割结果 + 模型

主干网络	方法	学习率调度	mIoU	#参数(M)	FLOPs(G)	配置	日志	模型检查点
MambaVision-T-1K	UPerNet	160K	46.0	55	945	配置	日志	模型
MambaVision-S-1K	UPerNet	160K	48.2	84	1135	配置	日志	模型
MambaVision-B-1K	UPerNet	160K	49.1	126	1342	配置	日志	模型
MambaVision-L3-512-21K	UPerNet	160K	53.2	780	3670	配置	日志	模型

安装

我们提供了一个 Docker 文件。此外，假设已经安装了最新版本的 PyTorch，可以通过运行以下命令来安装依赖项：

pip install -r requirements.txt

评估

可以使用以下命令在 ImageNet-1K 验证集上评估 MambaVision 模型：

python validate.py \
--model <模型名称>
--checkpoint <检查点路径>
--data_dir <ImageNet 路径>
--batch-size <每 GPU 的批量大小>

其中 --model 是 MambaVision 的变体（例如 mambavision_tiny_1k），--checkpoint 是预训练模型权重的路径，--data_dir 是 ImageNet-1K 验证集的路径，--batch-size 是批量大小。我们还提供了一个示例脚本 这里。

常见问题解答

1. MambaVision 是否支持处理任意输入分辨率的图像？

是的！您可以传递任意分辨率的图像，而无需更改模型。

2. 我有兴趣在我的仓库中重新实现 MambaVision。我们可以使用预训练权重吗？

是的！预训练权重以 CC-BY-NC-SA-4.0 许可发布。请在此仓库中提交一个问题，我们将把您的仓库添加到我们的代码库 README 中，并适当感谢您的努力。

3. 我可以将 MambaVision 应用于下游任务，如检测、分割吗？

是的！我们已经发布了支持下游任务的 模型，并提供了用于 目标检测 和 语义分割 的代码和预训练模型。

4. 如何计算每个模型的吞吐量和 FLOPs？

请参阅此 片段 以了解吞吐量和 FLOPs 的测量方法。结果可能因硬件而异。

引用

如果您发现 MambaVision 对您的工作有帮助，请考虑引用我们的论文：

@inproceedings{hatamizadeh2025mambavision,
  title={Mambavision: A hybrid mamba-transformer vision backbone},
  author={Hatamizadeh, Ali and Kautz, Jan},
  booktitle={Proceedings of the Computer Vision and Pattern Recognition Conference},
  pages={25261--25270},
  year={2025}
}

星标历史

许可证

版权所有 © 2026，英伟达公司。保留所有权利。

本作品根据 NVIDIA 源代码许可协议-NC 提供。点击 这里 查看该许可证的副本。

预训练模型以 CC-BY-NC-SA-4.0 许可共享。如果您对材料进行 remix、转换或构建，则必须以与原始材料相同的许可方式分发您的贡献。

有关 timm 仓库的许可证信息，请参阅其 仓库。

有关 ImageNet 数据集的许可证信息，请参阅 ImageNet 官方网站。

致谢

本仓库建立在 timm 仓库的基础上。我们感谢 罗斯·赖特曼 创建并维护这个高质量的库。

MambaVision 快速上手指南

MambaVision 是一种混合了 Mamba 与 Transformer 架构的视觉骨干网络，在 Top-1 准确率和吞吐量之间实现了新的最优平衡（SOTA Pareto-front）。本指南将帮助您快速在中国开发环境中部署并使用该模型。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 macOS
• Python: 3.8 或更高版本
• GPU: 支持 CUDA 的 NVIDIA GPU (推荐用于推理和训练)
• 核心依赖:
  • PyTorch >= 1.12
  • torchvision
  • timm (PyTorch Image Models)
  • transformers (Hugging Face)

建议操作： 国内用户建议使用清华源或阿里源加速 Python 包安装。

# 示例：配置 pip 使用清华源
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

您可以通过两种主要方式安装 MambaVision：使用官方 pip 包（推荐用于快速推理）或通过 Hugging Face Transformers 库。

方式一：通过 pip 安装（官方包）

这是最直接的集成方式，适合直接调用模型进行推理或微调。

pip install mambavision

如果需要安装额外的视觉处理依赖（如 timm），请确保已安装：

pip install timm pillow requests

方式二：通过 Hugging Face Transformers 使用

如果您习惯使用 transformers 生态，无需单独安装 mambavision 包，只需确保安装了最新版 transformers：

pip install --upgrade transformers

基本使用

以下提供两种最常用的使用场景代码示例。

场景 1：使用 pip 包进行图像分类

此方法适合直接使用官方预训练权重进行推理，支持任意分辨率输入。

from mambavision import create_model
import torch

# 1. 创建模型 (以 Tiny 版本为例)
# pretrained=True 会自动下载权重，也可指定本地 model_path
model = create_model('mamba_vision_T', pretrained=True)

# 2. 切换至评估模式并移至 GPU
model.cuda().eval()

# 3. 准备输入数据 (支持任意分辨率，此处示例为 512x224)
# 实际使用时请使用 transforms 对真实图片进行预处理 (归一化等)
dummy_image = torch.rand(1, 3, 512, 224).cuda()

# 4. 执行推理
with torch.no_grad():
    output = model(dummy_image)

# 输出维度为 [batch_size, 1000] (ImageNet 类别数)
print(f"Output logits shape: {output.shape}")

场景 2：使用 Hugging Face 进行特征提取

此方法适合需要获取多级特征图（用于检测、分割等下游任务）的场景。

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

# 1. 加载预训练模型
model = AutoModel.from_pretrained("nvidia/MambaVision-T-1K", trust_remote_code=True)
model.cuda().eval()

# 2. 准备图像数据
url = 'http://images.cocodataset.org/val2017/000000020247.jpg'
image = Image.open(requests.get(url, stream=True).raw)

# 定义输入分辨率 (MambaVision 支持动态分辨率)
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()

# 3. 执行推理并提取特征
with torch.no_grad():
    # out_avg_pool: 全局平均池化后的特征 (用于分类)
    # features: 包含 4 个阶段的多尺度特征列表 (用于检测/分割)
    out_avg_pool, features = model(inputs)

# 打印特征信息
print(f"全局池化特征尺寸：{out_avg_pool.size()}") 
print(f"特征阶段数量：{len(features)}")
print(f"第 1 阶段特征尺寸：{features[0].size()}")
print(f"第 4 阶段特征尺寸：{features[3].size()}")

可用模型列表

目前支持的预训练模型包括：

• ImageNet-1K: MambaVision-T, T2, S, B, L, L2
• ImageNet-21K: MambaVision-B, L, L2, L3 (高分辨率版本)

您可以在 Hugging Face Collection 查看所有模型权重。