新闻

• 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

模型权重

注意：

• ⁺ 表示我们以更细粒度进行短期微调。

对提供权重的评估

要在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}
}