VMamba

VMamba

高效的线性时间复杂度视觉骨干网络

VMamba是一种创新的视觉骨干网络,将Mamba状态空间语言模型应用于计算机视觉。其核心是视觉状态空间块堆栈,结合2D选择性扫描模块,实现线性时间复杂度。VMamba在图像分类、目标检测和语义分割等多项视觉任务中表现出色,特别是在输入尺度扩展效率方面优于现有模型。项目提供多种规模的预训练模型,适用于各类视觉感知任务。

VMamba计算机视觉深度学习图像处理神经网络Github开源项目
<div align="center"> <h1>VMamba</h1> <h3>VMamba: 视觉状态空间模型</h3>

刘悦<sup>1</sup>,田云杰<sup>1</sup>,赵宇中<sup>1</sup>, 于洪天<sup>1</sup>, 谢凌曦<sup>2</sup>, 王耀威<sup>3</sup>, 叶齐祥<sup>1</sup>, 刘云帆<sup>1</sup>

<sup>1</sup> 中国科学院大学, <sup>2</sup> 华为公司, <sup>3</sup> 鹏城实验室

论文: (arXiv 2401.10166)

</div>

: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作为计算机视觉的通用骨干网络。
<p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/6cf7cae3-3fe9-4657-b3ba-30ee19e6648d.png" alt="架构" width="80%"> </p>
  • VMamba的2D选择性扫描
<p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/3ed2c93c-cbfc-40e3-8a38-7663ada707e8.png" alt="架构" width="80%"> </p>
  • VMamba具有全局有效感受野
<p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/be395508-d025-4fa9-aedd-1fbfed79d530.png" alt="erf" width="80%"> </p>
  • VMamba在激活图上类似于基于Transformer的方法
<p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/12abeaa4-bf13-4d42-88c5-172f7b7b0283.png" alt="注意力" width="80%"> </p> <p align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/928b1f43-a2eb-48b6-abe2-09d66010202e.png" alt="激活" width="80%"> </p>

主要结果

:book: 详情请参阅performance.md

ImageNet-1K分类

名称预训练分辨率acc@1参数量FLOPs吞吐量训练吞吐量配置/日志/检查点
Swin-TImageNet-1K224x22481.228M4.5G1244987--
Swin-SImageNet-1K224x22483.250M8.7G718642--
Swin-BImageNet-1K224x22483.588M15.4G458496--
VMamba-S[s2l15]ImageNet-1K224x22483.650M8.7G877314配置/日志/检查点
VMamba-B[s2l15]ImageNet-1K224x22483.989M15.4G646247配置/日志/检查点
VMamba-T[s1l8]ImageNet-1K224x22482.630M4.9G1686571配置/日志/检查点
  • 本小节中的模型是从随机或手动初始化开始训练的。超参数继承自Swin,除了drop_path_rateEMA。除Vanilla-VMamba-T外,所有模型都使用EMA训练。
  • 吞吐量训练吞吐量是在配备AMD EPYC 7542 CPU的A100 GPU上评估的,批量大小为128。训练吞吐量使用混合分辨率测试,不包括优化器的时间消耗。
  • FLOPs参数量现在包括head(在之前的版本中,它们是不包括head计算的,所以数字略有增加)。
  • 我们使用@albertgu 提供的算法计算FLOPs,这将比之前的计算(基于selective_scan_ref函数,忽略了硬件感知算法)得到更大的数值。

COCO上的目标检测

骨干网络参数量FLOPs检测器bboxAPbboxAP50bboxAP75segmAPsegmAP50segmAP75配置/日志/检查点
Swin-T48M267GMaskRCNN@1x42.765.246.839.362.242.2--
Swin-S69M354GMaskRCNN@1x44.866.648.940.963.444.2--
Swin-B107M496GMaskRCNN@1x46.9----42.3------
VMamba-S[s2l15]70M384GMaskRCNN@1x48.770.053.443.767.347.0配置/日志/检查点
VMamba-B[s2l15]108M485GMaskRCNN@1x49.271.454.044.168.347.7配置/日志/检查点
VMamba-B[s2l15]108M485GMaskRCNN@1x[bs8]49.270.953.943.967.747.6配置/日志/检查点
VMamba-T[s1l8]50M271GMaskRCNN@1x47.369.352.042.766.445.9配置/日志/检查点
:---::---::---::---::---::---::---::---::---::---::---:
Swin-T48M267GMaskRCNN@3x46.068.150.341.665.144.9--
Swin-S69M354GMaskRCNN@3x48.269.852.843.267.046.1--
VMamba-S[s2l15]70M384GMaskRCNN@3x49.970.954.744.2068.247.7配置/日志/检查点
VMamba-T[s1l8]50M271GMaskRCNN@3x48.870.453.5043.767.447.0配置/日志/检查点
  • 本小节中的模型是从分类中训练的模型初始化的。
  • 我们现在使用@albertgu 提供的算法计算FLOPs,这将比之前的计算(基于selective_scan_ref函数,并忽略了硬件感知算法)更大。

ADE20K上的语义分割

骨干网络输入尺寸参数量FLOPs分割器mIoU(SS)mIoU(MS)配置/日志/日志(ms)/检查点
Swin-T512x51260M945GUperNet@160k44.445.8--
Swin-S512x51281M1039GUperNet@160k47.649.5--
Swin-B512x512121M1188GUperNet@160k48.149.7--
VMamba-S[s2l15]512x51282M1028GUperNet@160k50.651.2配置/日志/日志(ms)/检查点
VMamba-B[s2l15]512x512122M1170GUperNet@160k51.051.6配置/日志/日志(ms)/检查点
VMamba-T[s1l8]512x51262M949GUperNet@160k47.948.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_ssmselective_scan是否足够接近instead

  • 如果你对selective scan感兴趣,可以查看mambamamba-minimamba.pymamba-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

检测和分割

使用mmdetectionmmsegmentation进行评估:

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

在分割中使用--tta获取mIoU(ms)

使用mmdetectionmmsegmentation进行训练:

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

有关检测和分割任务的更多信息,请参阅mmdetectionmmsegmentation的手册。记得使用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历史

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%)

编辑推荐精选

AEE

AEE

AI Excel全自动制表工具

AEE 在线 AI 全自动 Excel 编辑器,提供智能录入、自动公式、数据整理、图表生成等功能,高效处理 Excel 任务,提升办公效率。支持自动高亮数据、批量计算、不规则数据录入,适用于企业、教育、金融等多场景。

UI-TARS-desktop

UI-TARS-desktop

基于 UI-TARS 视觉语言模型的桌面应用,可通过自然语言控制计算机进行多模态操作。

UI-TARS-desktop 是一款功能强大的桌面应用,基于 UI-TARS(视觉语言模型)构建。它具备自然语言控制、截图与视觉识别、精确的鼠标键盘控制等功能,支持跨平台使用(Windows/MacOS),能提供实时反馈和状态显示,且数据完全本地处理,保障隐私安全。该应用集成了多种大语言模型和搜索方式,还可进行文件系统操作。适用于需要智能交互和自动化任务的场景,如信息检索、文件管理等。其提供了详细的文档,包括快速启动、部署、贡献指南和 SDK 使用说明等,方便开发者使用和扩展。

Wan2.1

Wan2.1

开源且先进的大规模视频生成模型项目

Wan2.1 是一个开源且先进的大规模视频生成模型项目,支持文本到图像、文本到视频、图像到视频等多种生成任务。它具备丰富的配置选项,可调整分辨率、扩散步数等参数,还能对提示词进行增强。使用了多种先进技术和工具,在视频和图像生成领域具有广泛应用前景,适合研究人员和开发者使用。

爱图表

爱图表

全流程 AI 驱动的数据可视化工具,助力用户轻松创作高颜值图表

爱图表(aitubiao.com)就是AI图表,是由镝数科技推出的一款创新型智能数据可视化平台,专注于为用户提供便捷的图表生成、数据分析和报告撰写服务。爱图表是中国首个在图表场景接入DeepSeek的产品。通过接入前沿的DeepSeek系列AI模型,爱图表结合强大的数据处理能力与智能化功能,致力于帮助职场人士高效处理和表达数据,提升工作效率和报告质量。

Qwen2.5-VL

Qwen2.5-VL

一款强大的视觉语言模型,支持图像和视频输入

Qwen2.5-VL 是一款强大的视觉语言模型,支持图像和视频输入,可用于多种场景,如商品特点总结、图像文字识别等。项目提供了 OpenAI API 服务、Web UI 示例等部署方式,还包含了视觉处理工具,有助于开发者快速集成和使用,提升工作效率。

HunyuanVideo

HunyuanVideo

HunyuanVideo 是一个可基于文本生成高质量图像和视频的项目。

HunyuanVideo 是一个专注于文本到图像及视频生成的项目。它具备强大的视频生成能力,支持多种分辨率和视频长度选择,能根据用户输入的文本生成逼真的图像和视频。使用先进的技术架构和算法,可灵活调整生成参数,满足不同场景的需求,是文本生成图像视频领域的优质工具。

WebUI for Browser Use

WebUI for Browser Use

一个基于 Gradio 构建的 WebUI,支持与浏览器智能体进行便捷交互。

WebUI for Browser Use 是一个强大的项目,它集成了多种大型语言模型,支持自定义浏览器使用,具备持久化浏览器会话等功能。用户可以通过简洁友好的界面轻松控制浏览器智能体完成各类任务,无论是数据提取、网页导航还是表单填写等操作都能高效实现,有利于提高工作效率和获取信息的便捷性。该项目适合开发者、研究人员以及需要自动化浏览器操作的人群使用,在 SEO 优化方面,其关键词涵盖浏览器使用、WebUI、大型语言模型集成等,有助于提高网页在搜索引擎中的曝光度。

xiaozhi-esp32

xiaozhi-esp32

基于 ESP32 的小智 AI 开发项目,支持多种网络连接与协议,实现语音交互等功能。

xiaozhi-esp32 是一个极具创新性的基于 ESP32 的开发项目,专注于人工智能语音交互领域。项目涵盖了丰富的功能,如网络连接、OTA 升级、设备激活等,同时支持多种语言。无论是开发爱好者还是专业开发者,都能借助该项目快速搭建起高效的 AI 语音交互系统,为智能设备开发提供强大助力。

olmocr

olmocr

一个用于 OCR 的项目,支持多种模型和服务器进行 PDF 到 Markdown 的转换,并提供测试和报告功能。

olmocr 是一个专注于光学字符识别(OCR)的 Python 项目,由 Allen Institute for Artificial Intelligence 开发。它支持多种模型和服务器,如 vllm、sglang、OpenAI 等,可将 PDF 文件的页面转换为 Markdown 格式。项目还提供了测试框架和 HTML 报告生成功能,方便用户对 OCR 结果进行评估和分析。适用于科研、文档处理等领域,有助于提高工作效率和准确性。

飞书多维表格

飞书多维表格

飞书多维表格 ×DeepSeek R1 满血版

飞书多维表格联合 DeepSeek R1 模型,提供 AI 自动化解决方案,支持批量写作、数据分析、跨模态处理等功能,适用于电商、短视频、影视创作等场景,提升企业生产力与创作效率。关键词:飞书多维表格、DeepSeek R1、AI 自动化、批量处理、企业协同工具。

下拉加载更多