基于动态空间稀疏化的高效视觉Transformer和CNN

本仓库包含DynamicViT（NeurIPS 2021）的PyTorch实现。

DynamicViT 是一种动态标记稀疏化框架，能够根据输入逐步且动态地剪枝视觉Transformer中的冗余标记。我们的方法可将FLOPs降低超过 30%，并将吞吐量提升超过 40%，同时各类视觉Transformer的精度下降幅度不超过 0.5%。

[项目页面] [arXiv (NeurIPS 2021)]

🔥更新

我们已将该方法扩展至更多 网络架构（如ConvNeXt和Swin Transformer）以及更多 任务（如目标检测和语义分割），并引入了改进的 动态空间稀疏化 框架。详情请参阅论文的扩展版本，该版本已被T-PAMI接收。

[arXiv (T-PAMI, 期刊版)]

图像示例

---

视频示例

模型库

我们提供了在ImageNet上预训练的DynamicViT模型：

名称	模型	rho	精度@1	精度@5	FLOPs	下载链接
DynamicViT-DeiT-256/0.7	deit-256	0.7	76.53	93.12	1.3G	Google Drive / 清华大学云
DynamicViT-DeiT-S/0.7	deit-s	0.7	79.32	94.68	2.9G	Google Drive / 清华大学云
DynamicViT-DeiT-B/0.7	deit-b	0.7	81.43	95.46	11.4G	Google Drive / 清华大学云
DynamicViT-LVViT-S/0.5	lvvit-s	0.5	81.97	95.76	3.7G	Google Drive / 清华大学云
DynamicViT-LVViT-S/0.7	lvvit-s	0.7	83.08	96.25	4.6G	Google Drive / 清华大学云
DynamicViT-LVViT-M/0.7	lvvit-m	0.7	83.82	96.58	8.5G	Google Drive / 清华大学云

---

🔥更新： 我们还提供了在ImageNet上预训练的DynamicCNN和DynamicSwin模型：

名称	模型	rho	精度@1	精度@5	FLOPs	下载链接
DynamicCNN-T/0.7	convnext-t	0.7	81.59	95.72	3.6G	Google Drive / 清华大学云
DynamicCNN-T/0.9	convnext-t	0.9	82.06	95.89	3.9G	Google Drive / 清华大学云
DynamicCNN-S/0.7	convnext-s	0.7	82.57	96.29	5.8G	Google Drive / 清华大学云
DynamicCNN-S/0.9	convnext-s	0.9	83.12	96.42	6.8G	Google Drive / 清华大学云
DynamicCNN-B/0.7	convnext-b	0.7	83.45	96.56	10.2G	Google Drive / 清华大学云
DynamicCNN-B/0.9	convnext-b	0.9	83.96	96.76	11.9G	Google Drive / 清华大学云
DynamicSwin-T/0.7	swin-t	0.7	80.91	95.42	4.0G	Google Drive / 清华大学云
DynamicSwin-S/0.7	swin-s	0.7	83.21	96.33	6.9G	Google Drive / 清华大学云
DynamicSwin-B/0.7	swin-b	0.7	83.43	96.45	12.1G	Google Drive / 清华大学云

使用方法

需求

• torch>=1.8.0
• torchvision>=0.9.0
• timm==0.3.2
• tensorboardX
• six
• fvcore

数据准备: 从 http://image-net.org/ 下载并解压 ImageNet 图像。目录结构应如下所示：

│ILSVRC2012/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   ├── n01440764_10027.JPEG
│  │   ├── ......
│  ├── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   ├── ILSVRC2012_val_00002138.JPEG
│  │   ├── ......
│  ├── ......

模型准备: 如有需要，请下载预训练模型：

模型	链接	模型	链接
DeiT-Small	链接	LVViT-S	链接
DeiT-Base	链接	LVViT-M	链接
ConvNeXt-T	链接	Swin-T	链接
ConvNeXt-S	链接	Swin-S	链接
ConvNeXt-B	链接	Swin-B	链接

演示

您可以在 Colab 上试用 DynamicViT 。感谢 @dirtycomputer 的贡献。

我们还提供了一个 Jupyter 笔记本，您可以在其中运行 DynamicViT 的可视化效果。

要运行演示，您需要安装 matplotlib。

评估

要在单个 GPU 上使用 ImageNet 验证集评估预训练的 DynamicViT 模型，请运行以下命令：

python infer.py --data_path /path/to/ILSVRC2012/ --model model_name \
--model_path /path/to/model --base_rate 0.7 

训练

要在 ImageNet 上训练动态空间稀疏化模型，请运行以下命令：

（您可以通过调整 base_rate 来训练不同保留率的模型。）

DeiT-S

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --output_dir logs/dynamicvit_deit-s --model deit-s --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 30 --base_rate 0.7 --lr 1e-3 --warmup_epochs 5

DeiT-B

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --output_dir logs/dynamicvit_deit-b --model deit-b --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 30 --base_rate 0.7 --lr 1e-3 --warmup_epochs 5 --drop_path 0.2 --ratio_weight 5.0

LV-ViT-S

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --output_dir logs/dynamicvit_lvvit-s --model lvvit-s --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 30 --base_rate 0.7 --lr 1e-3 --warmup_epochs 5

LV-ViT-M

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --output_dir logs/dynamicvit_lvvit-m --model lvvit-m --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 30 --base_rate 0.7 --lr 1e-3 --warmup_epochs 5

DynamicViT 仅需训练 15 个 epoch 也能达到相近的性能（相比 30 个 epoch 的准确率低约 0.1%）。

ConvNeXt-T

在 8 张 GPU 上训练：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --output_dir logs/dynamic_conv-t --model convnext-t --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.2 --update_freq 4 --lr_scale 0.2

在 4 个 8-GPU 节点上训练：

python run_with_submitit.py --nodes 4 --ngpus 8 --output_dir logs/dynamic_conv-t --model convnext-t --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.2 --update_freq 1 --lr_scale 0.2

ConvNeXt-S

在 8 张 GPU 上训练：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --output_dir logs/dynamic_conv-s --model convnext-s --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.2 --update_freq 4 --lr_scale 0.2

在 4 个 8-GPU 节点上训练：

python run_with_submitit.py --nodes 4 --ngpus 8 --output_dir logs/dynamic_conv-s --model convnext-s --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.2 --update_freq 1 --lr_scale 0.2

ConvNeXt-B

在 8 张 GPU 上训练：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --output_dir logs/dynamic_conv-b --model convnext-b --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.5 --update_freq 4 --lr_scale 0.2

在 4 个 8-GPU 节点上训练：

python run_with_submitit.py --nodes 4 --ngpus 8 --output_dir logs/dynamic_conv-b --model convnext-b --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.5 --update_freq 1 --lr_scale 0.2

Swin-T

在 8 张 GPU 上训练：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --output_dir logs/dynamic_swin-t --model swin-t --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.2 --update_freq 4 --lr_scale 0.2

在 4 个 8-GPU 节点上训练：

python run_with_submitit.py --nodes 4 --ngpus 8 --output_dir logs/dynamic_swin-t --model swin-t --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.2 --update_freq 1 --lr_scale 0.2

Swin-S

在 8 张 GPU 上训练：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --output_dir logs/dynamic_swin-s --model swin-s --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.2 --update_freq 4 --lr_scale 0.2

在 4 个 8-GPU 节点上训练：

python run_with_submitit.py --nodes 4 --ngpus 8 --output_dir logs/dynamic_swin-s --model swin-s --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.2 --update_freq 1 --lr_scale 0.2

Swin-B

在 8 张 GPU 上训练：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --output_dir logs/dynamic_swin-b --model swin-b --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base_rate 0.7 --lr 4e-3 --drop_path 0.5 --update_freq 4 --lr_scale 0.2

在 4 个 8-GPU 节点上训练：

python run_with_submitit.py --nodes 4 --ngpus 8 --output_dir logs/dynamic_swin-b --model swin-b --input_size 224 --batch_size 128 --data_path /path/to/ILSVRC2012/ --epochs 120 --base rate 0.7 --lr 4e-3 --drop_path 0.5 --update freq 1 --lr scale 0.2

许可证

MIT 许可证

致谢

我们的代码基于 pytorch-image-models、DeiT、LV-ViT、ConvNeXt 和 Swin-Transformer。

引用

如果您在研究中使用了我们的工作，请考虑引用以下文献：

@inproceedings{rao2021dynamicvit,
  title={DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification},
  author={Rao, Yongming and Zhao, Wenliang and Liu, Benlin and Lu, Jiwen and Zhou, Jie and Hsieh, Cho-Jui},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year = {2021}
}

@article{rao2022dynamicvit,
  title={Dynamic Spatial Sparsification for Efficient Vision Transformers and Convolutional Neural Networks},
  author={Rao, Yongming and Liu, Zuyan and Zhao, Wenliang and Zhou, Jie and Lu, Jiwen},
  journal={arXiv preprint arXiv:2207.01580},
  year={2022}

DynamicViT 快速上手指南

DynamicViT 是一个基于动态空间稀疏化（Dynamic Spatial Sparsification）的框架，旨在逐步且动态地剪枝视觉 Transformer（ViT）和 CNN 中的冗余 Token。该方法可在精度损失小于 0.5% 的前提下，降低超过 30% 的 FLOPs 并提升 40% 以上的吞吐量。支持 DeiT、LV-ViT、ConvNeXt 和 Swin Transformer 等多种架构。

环境准备

系统要求

• Python 3.x
• GPU 环境（推荐用于训练和推理）

前置依赖

请确保安装以下核心依赖库：

torch>=1.8.0
torchvision>=0.9.0
timm==0.3.2
tensorboardX
six
fvcore
matplotlib (仅用于可视化 Demo)

安装步骤

1. 克隆仓库

   git clone https://github.com/raoyongming/DynamicViT.git
   cd DynamicViT
   

2. 安装 Python 依赖 建议使用国内镜像源加速安装：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   注：若仓库无 requirements.txt 文件，请手动执行：

   pip install torch torchvision timm==0.3.2 tensorboardX six fvcore matplotlib -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 数据准备 下载 ImageNet 数据集 (ILSVRC2012)，目录结构需如下：

   ILSVRC2012/
   ├── train/
   │   ├── n01440764/
   │   └── ...
   └── val/
       ├── n01440764/
       └── ...
   

4. 模型准备（可选） 如需使用预训练模型，可从下表下载。国内用户推荐使用 清华云盘 链接：

   模型架构	推荐下载源 (清华云盘)
   DynamicViT-DeiT 系列	点击下载
   DynamicCNN (ConvNeXt)	点击下载
   DynamicSwin	点击下载

   完整模型列表及对应保持率 (rho) 请参考项目原文档 Model Zoo 部分。

基本使用

1. 快速体验 (Demo)

项目提供了 Colab 在线演示和本地 Jupyter Notebook 可视化示例。

• 在线运行: 点击 Colab Demo 直接在浏览器中体验。
• 本地可视化: 运行以下命令查看 Token 剪枝效果（需安装 matplotlib）：

  jupyter notebook viz_example.ipynb
  

2. 模型推理 (Evaluation)

使用单张 GPU 在 ImageNet 验证集上评估预训练模型。以下以 deit-s 模型为例，设置保留率 (base_rate) 为 0.7：

python infer.py --data_path /path/to/ILSVRC2012/ --model deit-s \
--model_path /path/to/pretrained_model.pth --base_rate 0.7

请将 /path/to/... 替换为实际的数据集路径和模型权重文件路径。

3. 模型训练 (Training)

基于 ImageNet 训练动态稀疏化模型。以下以 DeiT-Small 为例，在 8 卡环境下训练 30 个 epoch：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py \
--output_dir logs/dynamicvit_deit-s \
--model deit-s \
--input_size 224 \
--batch_size 128 \
--data_path /path/to/ILSVRC2012/ \
--epochs 30 \
--base_rate 0.7 \
--lr 1e-3 \
--warmup_epochs 5

提示：

• 可通过调整 --base_rate 参数控制 Token 保留比例（如 0.5, 0.7, 0.9）。
• 对于 ConvNeXt 或 Swin 等较大模型，建议参考原文档调整 drop_path、update_freq 及学习率策略。