Vim

<div align="center"> <h1>视觉曼巴</h1> <h3>基于双向状态空间模型的高效视觉表征学习</h3>

朱良辉<sup>1</sup> *,廖本澄<sup>1</sup> *,张茜<sup>2</sup>, 王鑫龙<sup>3</sup>, 刘文预<sup>1</sup>, 王兴刚<sup>1 :email:</sup>

<sup>1</sup> 华中科技大学, <sup>2</sup> 地平线机器人, <sup>3</sup> 北京智源人工智能研究院

(*) 同等贡献, (<sup>:email:</sup>) 通讯作者。

ArXiv 预印本 (arXiv 2401.09417), HuggingFace 页面 (🤗 2401.09417)

</div>

新闻

2024年2月10日: 我们更新了Vim-tiny/small的权重和训练脚本。通过将类别标记放置在中间位置，Vim取得了更好的结果。更多细节可以在代码和我们更新的arXiv中找到。
2024年1月18日: 我们在Arxiv上发布了论文。代码/模型即将推出。敬请期待！☕️

摘要

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

概述

预训练环境

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
需求: 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 准确率	Hugginface 仓库
Vim-tiny	7M	76.1	93.0	https://huggingface.co/hustvl/Vim-tiny-midclstok
Vim-tiny<sup>+</sup>	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<sup>+</sup>	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

注意：

<sup>+</sup> 表示我们以更细粒度进行短期微调。

对提供权重的评估

要在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条目引用它。

 @article{vim,
  title={Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model},
  author={Lianghui Zhu and Bencheng Liao and Qian Zhang and Xinlong Wang and Wenyu Liu and Xinggang Wang},
  journal={arXiv preprint arXiv:2401.09417},
  year={2024}
}