🔥 仅需一个文件，以最少步骤使用VMamba！

conda create -n vmamba python=3.10
pip install torch==2.2 torchvision torchaudio triton pytest chardet yacs termcolor fvcore seaborn packaging ninja einops numpy==1.24.4 timm==0.4.12
pip install https://github.com/state-spaces/mamba/releases/download/v2.2.4/mamba_ssm-2.2.4+cu12torch2.2cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
python vmamba.py

• 更新
• 摘要
• 概述
• 主要结果
• 快速入门
• 星标历史
• 引用
• 致谢

:white_check_mark: 更新

• 2024年9月25日：更新：VMamba已被NeurIPS 2024接收（亮点论文）！
• 2024年6月14日：更新：我们优化了代码，使其更易读；增加了对mamba2的支持。
• 2024年5月26日：更新：我们发布了VMambav2的更新权重，并附带新的arXiv论文。
• 2024年5月7日：更新：重要提示！ 在下游任务中使用torch.backends.cudnn.enabled=True可能会导致运行速度较慢。如果您发现VMamba在您的机器上运行较慢，请在vmamba.py中将其禁用，否则无需理会。
• ...

详情请参阅detailed_updates.md

摘要

设计计算效率高的网络架构仍然是计算机视觉领域持续的需求。本文将状态空间语言模型Mamba移植到VMamba中，构建了一种线性时间复杂度的视觉骨干网络。VMamba的核心是由带有二维选择性扫描（SS2D）模块的视觉状态空间（VSS）块堆叠而成。通过沿四条扫描路径遍历，SS2D有助于弥合一维选择性扫描的有序特性与二维视觉数据非序列结构之间的差距，从而促进从不同来源和视角收集上下文信息。基于VSS块，我们开发了一系列VMamba架构，并通过一系列架构和实现上的改进进一步加速其性能。大量实验表明，VMamba在多种视觉感知任务中表现出色，尤其在输入尺度扩展效率方面优于现有基准模型。

概述

• VMamba 是一种通用的计算机视觉骨干网络。

• VMamba的二维选择性扫描

• VMamba具有全局有效感受野

• VMamba在激活图上类似于基于Transformer的方法

主要结果

:book: 详情请参阅performance.md。

ImageNet-1K上的分类

名称	预训练	分辨率	acc@1	参数量	FLOPs	TP.	训练TP.	配置/日志/检查点
Swin-T	ImageNet-1K	224x224	81.2	28M	4.5G	1244	987	--
Swin-S	ImageNet-1K	224x224	83.2	50M	8.7G	718	642	--
Swin-B	ImageNet-1K	224x224	83.5	88M	15.4G	458	496	--
VMamba-S[s2l15]	ImageNet-1K	224x224	83.6	50M	8.7G	877	314	配置/日志/检查点
VMamba-B[s2l15]	ImageNet-1K	224x224	83.9	89M	15.4G	646	247	配置/日志/检查点
VMamba-T[s1l8]	ImageNet-1K	224x224	82.6	30M	4.9G	1686	571	配置/日志/检查点

• 本小节中的模型均采用随机或手动初始化从头开始训练。超参数沿用了Swin的设置，但drop_path_rate和EMA除外。除Vanilla-VMamba-T外，所有模型均使用EMA进行训练。
• TP.(吞吐量)和Train TP. (训练吞吐量)是在A100 GPU搭配AMD EPYC 7542 CPU的环境下，以批量大小128进行评估的。训练吞吐量测试时采用了混合分辨率，未计入优化器的时间消耗。
• FLOPs和参数量现在包含了头部部分（在之前的版本中，这些数值并未计算头部，因此会略高一些）。
• 我们使用@albertgu提供的算法来计算FLOPs，该算法得出的结果会比之前基于selective_scan_ref函数且未考虑硬件优化的计算方法更大。

COCO 数据集上的目标检测

主干网络	参数量	FLOPs	检测器	bboxAP	bboxAP50	bboxAP75	segmAP	segmAP50	segmAP75	配置/日志/检查点
Swin-T	48M	267G	MaskRCNN@1x	42.7	65.2	46.8	39.3	62.2	42.2	--
Swin-S	69M	354G	MaskRCNN@1x	44.8	66.6	48.9	40.9	63.4	44.2	--
Swin-B	107M	496G	MaskRCNN@1x	46.9	--	--	42.3	--	--	--
VMamba-S[s2l15]	70M	384G	MaskRCNN@1x	48.7	70.0	53.4	43.7	67.3	47.0	配置/日志/检查点
VMamba-B[s2l15]	108M	485G	MaskRCNN@1x	49.2	71.4	54.0	44.1	68.3	47.7	配置/日志/检查点
VMamba-B[s2l15]	108M	485G	MaskRCNN@1x[bs8]	49.2	70.9	53.9	43.9	67.7	47.6	配置/日志/检查点
VMamba-T[s1l8]	50M	271G	MaskRCNN@1x	47.3	69.3	52.0	42.7	66.4	45.9	配置/日志/检查点
:---:	:---:	:---:	:---:	:---:	:---:	:---:	:---:	:---:	:---:	:---:
Swin-T	48M	267G	MaskRCNN@3x	46.0	68.1	50.3	41.6	65.1	44.9	--
Swin-S	69M	354G	MaskRCNN@3x	48.2	69.8	52.8	43.2	67.0	46.1	--
VMamba-S[s2l15]	70M	384G	MaskRCNN@3x	49.9	70.9	54.7	44.20	68.2	47.7	配置/日志/检查点
VMamba-T[s1l8]	50M	271G	MaskRCNN@3x	48.8	70.4	53.50	43.7	67.4	47.0	配置/日志/检查点

• 本小节中的模型均从分类任务中训练好的模型初始化。
• 我们现在使用 @albertgu 提供的算法计算 FLOPs（见 GitHub issue #110），这比之前的计算结果要大（之前的计算基于 selective_scan_ref 函数，未考虑硬件相关的优化算法）。

ADE20K 数据集上的语义分割

主干网络	输入尺寸	参数量	FLOPs	分割器	mIoU(SS)	mIoU(MS)	配置/日志/多尺度日志/检查点
Swin-T	512x512	60M	945G	UperNet@160k	44.4	45.8	--
Swin-S	512x512	81M	1039G	UperNet@160k	47.6	49.5	--
Swin-B	512x512	121M	1188G	UperNet@160k	48.1	49.7	--
VMamba-S[s2l15]	512x512	82M	1028G	UperNet@160k	50.6	51.2	配置/日志/多尺度日志/检查点
VMamba-B[s2l15]	512x512	122M	1170G	UperNet@160k	51.0	51.6	配置/日志/多尺度日志/检查点
VMamba-T[s1l8]	512x512	62M	949G	UperNet@160k	47.9	48.8	配置/日志/多尺度日志/检查点

• 本小节中的模型均从分类任务中训练好的模型初始化。
• 我们现在使用 @albertgu 提供的算法计算 FLOPs（见 GitHub issue #110），这比之前的计算结果要大（之前的计算基于 selective_scan_ref 函数，未考虑硬件相关的优化算法）。

入门

安装

步骤 1：克隆 VMamba 仓库：

首先，克隆 VMamba 仓库并进入项目目录：

git clone https://github.com/MzeroMiko/VMamba.git
cd VMamba

步骤 2：环境设置：

VMamba 建议使用 conda 创建虚拟环境，并通过 pip 安装依赖。请使用以下命令设置环境： 此外，我们推荐使用 PyTorch ≥ 2.0 和 CUDA ≥ 11.8，但较低版本的 PyTorch 和 CUDA 也受支持。

创建并激活一个新的 conda 环境

conda create -n vmamba
conda activate vmamba

安装依赖

pip install -r requirements.txt
cd kernels/selective_scan && pip install .

检查选择性扫描（可选）

• 如果您想将模块与 mamba_ssm 进行对比测试，请先安装 mamba_ssm！

• 若要检查我们的 selective scan 实现是否与 mamba_ssm 相同，可以运行 selective_scan/test_selective_scan.py。将 selective_scan/test_selective_scan.py 中的 MODE = "mamba_ssm_sscore" 修改为相应值，然后运行 pytest selective_scan/test_selective_scan.py。

• 若想验证我们的 selective scan 实现是否与参考代码（selective_scan_ref）一致，可将 selective_scan/test_selective_scan.py 中的 MODE = "sscore" 修改为相应值，再运行 pytest selective_scan/test_selective_scan.py。

• MODE = "mamba_ssm" 表示检查 mamba_ssm 的结果是否接近 selective_scan_ref，而 "sstest" 则保留用于开发。

• 如果发现 mamba_ssm（selective_scan_cuda）或 selective_scan（selctive_scan_cuda_core）与 selective_scan_ref 差距较大，测试失败也不要担心。请查看 mamba 是否已解决该问题。

• 如果您对选择性扫描感兴趣，可以进一步了解 mamba、mamba-mini、mamba.py 和 mamba-minimal 等项目。

检测与分割任务的依赖（可选）

pip install mmengine==0.10.1 mmcv==2.1.0 opencv-python-headless ftfy regex
pip install mmdet==3.3.0 mmsegmentation==1.2.2 mmpretrain==1.2.0

模型训练与推理

分类

要在 ImageNet 数据集上训练 VMamba 分类模型，可根据不同配置使用以下命令：

python -m torch.distributed.launch --nnodes=1 --node_rank=0 --nproc_per_node=8 --master_addr="127.0.0.1" --master_port=29501 main.py --cfg </path/to/config> --batch-size 128 --data-path </path/of/dataset> --output /tmp

若仅需测试性能（包括参数量和 FLOPs）：

python -m torch.distributed.launch --nnodes=1 --node_rank=0 --nproc_per_node=1 --master_addr="127.0.0.1" --master_port=29501 main.py --cfg </path/to/config> --batch-size 128 --data-path </path/of/dataset> --output /tmp --pretrained </path/of/checkpoint>

更多详情请参阅 modelcard。

检测与分割

使用 mmdetection 或 mmsegmentation 进行评估：

bash ./tools/dist_test.sh </path/to/config> </path/to/checkpoint> 1

• 使用 --tta 可获得分割任务中的 mIoU(ms)。

使用 mmdetection 或 mmsegmentation 进行训练：

bash ./tools/dist_train.sh </path/to/config> 8

有关检测和分割任务的更多信息，请参阅 mmdetection 和 mmsegmentation 的官方文档。请务必在 configs 目录中使用合适的骨干网络配置。

分析工具

VMamba 提供了用于可视化 Mamba “注意力”和有效感受野、分析吞吐量及训练吞吐量的工具。请使用以下命令进行分析：

# 可视化 Mamba “注意力”
CUDA_VISIBLE_DEVICES=0 python analyze/attnmap.py

# 分析有效感受野
CUDA_VISIBLE_DEVICES=0 python analyze/erf.py

# 分析吞吐量和训练吞吐量
CUDA_VISIBLE_DEVICES=0 python analyze/tp.py

我们还包含了一些在此项目中可能用到的其他分析工具，感谢所有为这些工具做出贡献的人。

星标历史

引用

@article{liu2024vmamba,
  title={VMamba: 视觉状态空间模型},
  author={刘悦、田云杰、赵宇中、于洪天、谢凌曦、王耀伟、叶启祥、刘云帆},
  journal={arXiv 预印本 arXiv:2401.10166},
  year={2024}
}

致谢

本项目基于 Mamba（论文、代码）、Swin-Transformer（论文、代码）、ConvNeXt（论文、代码）、OpenMMLab， 以及从 replknet 借鉴的 analyze/get_erf.py，感谢他们的杰出工作。

• 我们最近发布了 Fast-iTPN，据我们所知，该模型在 Tiny/Small/Base 级别上取得了 ImageNet-1K 数据集的最佳性能。（Tiny-24M-86.5%，Small-40M-87.8%，Base-85M-88.75%）

VMamba 快速上手指南

VMamba 是一种基于视觉状态空间模型（Visual State Space Model）的通用视觉骨干网络，具有线性时间复杂度，在图像分类、目标检测和语义分割等任务中表现优异。

环境准备

• 操作系统: Linux (x86_64)
• Python: 3.10
• GPU: 支持 CUDA 的 NVIDIA 显卡 (示例命令基于 CUDA 12)
• 编译器: 需安装 ninja 以加速编译

安装步骤

请按照以下顺序执行命令来配置环境并安装依赖。

1. 创建 Conda 环境

conda create -n vmamba python=3.10
conda activate vmamba

2. 安装基础依赖

推荐使用国内镜像源（如清华源）加速下载：

pip install torch==2.2 torchvision torchaudio triton pytest chardet yacs termcolor fvcore seaborn packaging ninja einops numpy==1.24.4 timm==0.4.12 -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 安装 Mamba SSM 核心组件

安装预编译的 mamba_ssm wheel 包（适用于 CUDA 12 + Torch 2.2 + Python 3.10）：

pip install https://github.com/state-spaces/mamba/releases/download/v2.2.4/mamba_ssm-2.2.4+cu12torch2.2cxx11abiTRUE-cp310-cp310-linux_x86_64.whl

注意：如果您的 CUDA 版本或 Python 版本不同，请访问 Mamba Releases 查找对应的 .whl 文件进行安装。

4. 获取项目代码

git clone https://github.com/MzeroMiko/VMamba.git
cd VMamba

基本使用

项目提供了一个极简的单文件测试脚本，用于验证安装是否成功并运行模型。

运行测试脚本

在项目根目录下直接运行：

python vmamba.py

性能优化提示

如果在下游任务中发现推理速度较慢，请检查 vmamba.py 中的设置。根据官方更新日志，在某些机器上禁用 cuDNN 可能提升速度： 确保代码中未强制开启 torch.backends.cudnn.enabled=True，或者根据实际测试结果手动调整该选项。

---

更多详细配置（如分类、检测、分割任务的训练与评估）请参考项目目录下的 classification, detection, segmentation 文件夹及对应的配置文件。