新闻

• 2024年5月2日: Vision Mamba (Vim) 被 ICML2024 接收。🎉 会议页面可在此查看 这里。

• 2024年2月10日: 我们更新了 Vim-tiny/small 的权重和训练脚本。通过将类别标记放置在中间位置，Vim 取得了更好的效果。更多细节请参阅代码及我们更新后的 arXiv。

• 2024年1月18日: 我们已在 ArXiv 上发布了论文。代码/模型即将发布，请持续关注！☕️

摘要

最近，具有高效硬件感知设计的状态空间模型（SSMs），即 Mamba 深度学习模型，已在长序列建模方面展现出巨大潜力。与此同时，完全基于 SSMs 构建高效且通用的视觉骨干网络是一个极具吸引力的方向。然而，由于视觉数据对位置敏感以及视觉理解需要全局上下文，SSMs 在表示视觉数据方面仍面临挑战。在本文中，我们证明了视觉表征学习并不依赖于自注意力机制，并提出了一种新的通用视觉骨干网络——双向 Mamba 块（Vim）。该网络通过位置嵌入标记图像序列，并利用双向状态空间模型压缩视觉表征。在 ImageNet 分类、COCO 目标检测和 ADE20k 语义分割任务上，Vim 的性能优于 DeiT 等成熟的视觉 Transformer 模型，同时计算和内存效率也显著提升。例如，在分辨率为 1248×1248 的图像上进行批量推理提取特征时，Vim 的速度比 DeiT 快 2.8 倍，且 GPU 内存占用减少了 86.8%。这些结果表明，Vim 能够克服高分辨率图像上执行 Transformer 式理解时的计算和内存限制，并有望成为下一代视觉基础模型的骨干网络。

概述

预训练环境

• NVIDIA GPU：

  • Python 3.10.13

    • conda create -n your_env_name python=3.10.13

  • torch 2.1.1 + cu118

    • pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118

• AMD GPU：

  • 建议使用 Docker 镜像，例如 rocm/pytorch:rocm6.2_ubuntu20.04_py3.9_pytorch_release_2.1.2。有关分步说明，请参阅 Vision Mamba 在 AMD GPU 上使用 ROCm 的指南。

• 依赖项：vim_requirements.txt

  • pip install -r vim/vim_requirements.txt

• 安装 causal_conv1d 和 mamba

  • pip install -e causal_conv1d>=1.1.0
  • pip install -e mamba-1p1p1

训练你的 Vim

bash vim/scripts/pt-vim-t.sh

以更细粒度训练你的 Vim

bash vim/scripts/ft-vim-t.sh

模型权重

模型	参数量	Top-1 准确率	Top-5 准确率	HuggingFace 仓库
Vim-tiny	7M	76.1	93.0	https://huggingface.co/hustvl/Vim-tiny-midclstok
Vim-tiny+	7M	78.3	94.2	https://huggingface.co/hustvl/Vim-tiny-midclstok
Vim-small	26M	80.5	95.1	https://huggingface.co/hustvl/Vim-small-midclstok
Vim-small+	26M	81.6	95.4	https://huggingface.co/hustvl/Vim-small-midclstok
Vim-base	98M	81.9	95.8	https://huggingface.co/hustvl/Vim-base-midclstok

注释：

• ⁺ 表示我们在较短的时间内以更细粒度进行微调。

使用提供的权重进行评估

要在 ImageNet-1K 数据集上评估 Vim-Ti，请运行：

python main.py --eval --resume /path/to/ckpt --model vim_tiny_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2 --data-path /path/to/imagenet

致谢 :heart:

本项目基于 Mamba (论文, 代码)、Causal-Conv1d (代码)、DeiT (论文, 代码)。感谢他们的杰出工作。

引用

如果您在研究或应用中发现 Vim 有用，请考虑给我们一颗星 🌟，并使用以下 BibTeX 条目引用：

 @inproceedings{vim,
  title={Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model},
  author={Zhu, Lianghui and Liao, Bencheng and Zhang, Qian and Wang, Xinlong and Liu, Wenyu and Wang, Xinggang},
  booktitle={Forty-first International Conference on Machine Learning}
}

Vision Mamba (Vim) 快速上手指南

Vision Mamba (Vim) 是一种基于双向状态空间模型（Bidirectional SSM）的高效视觉骨干网络。它在 ImageNet 分类、COCO 目标检测和 ADE20k 语义分割等任务上，相比 DeiT 等主流 Vision Transformer 模型，在保持更高性能的同时，显著提升了计算速度和内存效率。

环境准备

系统要求

• GPU: NVIDIA GPU (推荐) 或 AMD GPU
• Python: 3.10.13
• CUDA: 11.8 (针对 NVIDIA 用户)

前置依赖

本项目依赖 torch, causal_conv1d, 和 mamba 等核心库。

安装步骤

1. 创建 Conda 环境

首先创建一个名为 vim_env 的虚拟环境（名称可自定义），并指定 Python 版本：

conda create -n vim_env python=3.10.13
conda activate vim_env

2. 安装 PyTorch

NVIDIA GPU 用户： 安装与 CUDA 11.8 兼容的 PyTorch 2.1.1 版本：

pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118

AMD GPU 用户： 建议使用官方提供的 Docker 镜像（如 rocm/pytorch:rocm6.2_ubuntu20.04_py3.9_pytorch_release_2.1.2）。详细步骤请参考 Vision Mamba on AMD GPU with ROCm。

3. 安装项目依赖

克隆项目代码后，进入目录安装基础依赖：

pip install -r vim/vim_requirements.txt

4. 安装核心算子

依次安装 causal_conv1d 和 mamba 扩展库：

pip install -e causal_conv1d>=1.1.0
pip install -e mamba-1p1p1

基本使用

模型权重下载

你可以从 HuggingFace 下载预训练权重。以下是部分可用模型：

模型	参数量	Top-1 准确率	HuggingFace 地址
Vim-tiny	7M	76.1%	hustvl/Vim-tiny-midclstok
Vim-small	26M	80.5%	hustvl/Vim-small-midclstok
Vim-base	98M	81.9%	hustvl/Vim-base-midclstok

> 注：带 + 号的模型表示经过更细粒度的微调，性能更优。

推理评估示例

以下命令演示如何在 ImageNet-1K 数据集上评估 Vim-Tiny 模型。请替换 <path/to/ckpt> 为你的权重文件路径，<path/to/imagenet> 为数据集路径。

python main.py --eval --resume /path/to/ckpt --model vim_tiny_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2 --data-path /path/to/imagenet

训练示例

如果你希望从头开始训练或使用提供的脚本进行微调，可以运行以下命令：

预训练 (Pretraining):

bash vim/scripts/pt-vim-t.sh

细粒度微调 (Finer Granularity Finetuning):

bash vim/scripts/ft-vim-t.sh