VMamba

VMamba: 视觉状态空间模型

刘悦¹,田云杰¹,赵宇中¹, 于洪天¹, 谢凌曦², 王耀威³, 叶齐祥¹, 刘云帆¹

¹ 中国科学院大学, ² 华为公司, ³ 鹏城实验室

论文: (arXiv 2401.10166)

更新
摘要
概述
主要结果
快速开始
Star历史
引用
致谢

:white_check_mark: 更新

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的核心是一堆带有2D选择性扫描（SS2D）模块的视觉状态空间（VSS）块。通过沿着四条扫描路线遍历，SS2D帮助弥合了1D选择性扫描的有序性质与2D视觉数据的非序列结构之间的差距，这有助于从各种来源和角度收集上下文信息。基于VSS块，我们开发了一系列VMamba架构，并通过一系列架构和实现改进来加速它们。大量实验展示了VMamba在各种视觉感知任务中的出色性能，突出了其相比现有基准模型在输入规模效率方面的优势。

概述

VMamba作为计算机视觉的通用骨干网络。

VMamba的2D选择性扫描

VMamba具有全局有效感受野

erf

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

注意力

主要结果

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

ImageNet-1K分类

名称	预训练	分辨率	acc@1	参数量	FLOPs	吞吐量	训练吞吐量	配置/日志/检查点
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训练。
吞吐量和训练吞吐量是在配备AMD EPYC 7542 CPU的A100 GPU上评估的，批量大小为128。训练吞吐量使用混合分辨率测试，不包括优化器的时间消耗。
FLOPs和参数量现在包括head（在之前的版本中，它们是不包括head计算的，所以数字略有增加）。
我们使用@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，这将比之前的计算（基于selective_scan_ref函数，并忽略了硬件感知算法）更大。

ADE20K上的语义分割

骨干网络	输入尺寸	参数量	FLOPs	分割器	mIoU(SS)	mIoU(MS)	配置/日志/日志(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	配置/日志/日志(ms)/检查点
VMamba-B[`s2l15`]	512x512	122M	1170G	UperNet@160k	51.0	51.6	配置/日志/日志(ms)/检查点
VMamba-T[`s1l8`]	512x512	62M	949G	UperNet@160k	47.9	48.8	配置/日志/日志(ms)/检查点

本小节中的模型是从分类中训练的模型初始化的。
我们现在使用@albertgu 提供的算法计算FLOPs，这将比之前的计算（基于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。在该文件中将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_ssm和selective_scan是否足够接近instead。
如果你对selective scan感兴趣,可以查看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

如果你只想测试性能(同时包括参数和浮点运算):

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

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

Star历史

引用

@article{liu2024vmamba,
  title={VMamba: Visual State Space Model},
  author={Liu, Yue and Tian, Yunjie and Zhao, Yuzhong and Yu, Hongtian and Xie, Lingxi and Wang, Yaowei and Ye, Qixiang and Liu, Yunfan},
  journal={arXiv preprint arXiv:2401.10166},
  year={2024}
}

致谢

本项目基于Mamba (论文, 代码), Swin-Transformer (论文, 代码), ConvNeXt (论文, 代码), OpenMMLab, analyze/get_erf.py采用自replknet,感谢他们的优秀工作。

我们最近发布了Fast-iTPN,据我们所知,它在ImageNet-1K的Tiny/Small/Base级别模型上报告了最佳性能。(Tiny-24M-86.5%, Small-40M-87.8%, Base-85M-88.75%)