This repository contains PyTorch implementation for DynamicViT (NeurIPS 2021).
DynamicViT is a dynamic token sparsification framework to prune redundant tokens in vision transformers progressively and dynamically based on the input. Our method can reduces over 30% FLOPs and improves the throughput by over 40% while the drop of accuracy is within 0.5% for various vision transformers.
[Project Page] [arXiv (NeurIPS 2021)]
We extend our method to more network architectures (i.e., ConvNeXt and Swin Transformers) and more tasks (i.e., object detection and semantic segmentation) with an improved dynamic spatial sparsification framework. Please refer to the extended version of our paper for details. The extended version has been accepted by T-PAMI.
[arXiv (T-PAMI, Journal Version)]
We provide our DynamicViT models pretrained on ImageNet:
| name | model | rho | acc@1 | acc@5 | FLOPs | url |
|---|---|---|---|---|---|---|
| DynamicViT-DeiT-256/0.7 | deit-256 |
0.7 | 76.53 | 93.12 | 1.3G | Google Drive / Tsinghua Cloud |
| DynamicViT-DeiT-S/0.7 | deit-s |
0.7 | 79.32 | 94.68 | 2.9G | Google Drive / Tsinghua Cloud |
| DynamicViT-DeiT-B/0.7 | deit-b |
0.7 | 81.43 | 95.46 | 11.4G | Google Drive / Tsinghua Cloud |
| DynamicViT-LVViT-S/0.5 | lvvit-s |
0.5 | 81.97 | 95.76 | 3.7G | Google Drive / Tsinghua Cloud |
| DynamicViT-LVViT-S/0.7 | lvvit-s |
0.7 | 83.08 | 96.25 | 4.6G | Google Drive / Tsinghua Cloud |
| DynamicViT-LVViT-M/0.7 | lvvit-m |
0.7 | 83.82 | 96.58 | 8.5G | Google Drive / Tsinghua Cloud |
🔥Updates: We provide our DynamicCNN and DynamicSwin models pretrained on ImageNet:
| name | model | rho | acc@1 | acc@5 | FLOPs | url |
|---|---|---|---|---|---|---|
| DynamicCNN-T/0.7 | convnext-t |
0.7 | 81.59 | 95.72 | 3.6G | Google Drive / Tsinghua Cloud |
| DynamicCNN-T/0.9 | convnext-t |
0.9 | 82.06 | 95.89 | 3.9G | Google Drive / Tsinghua Cloud |
| DynamicCNN-S/0.7 | convnext-s |
0.7 | 82.57 | 96.29 | 5.8G | Google Drive / Tsinghua Cloud |
| DynamicCNN-S/0.9 | convnext-s |
0.9 | 83.12 | 96.42 | 6.8G | Google Drive / Tsinghua Cloud |
| DynamicCNN-B/0.7 | convnext-b |
0.7 | 83.45 | 96.56 | 10.2G | Google Drive / Tsinghua Cloud |
| DynamicCNN-B/0.9 | convnext-b |
0.9 | 83.96 | 96.76 | 11.9G | Google Drive / Tsinghua Cloud |
| DynamicSwin-T/0.7 | swin-t |
0.7 | 80.91 | 95.42 | 4.0G | Google Drive / Tsinghua Cloud |
| DynamicSwin-S/0.7 | swin-s |
0.7 | 83.21 | 96.33 | 6.9G | Google Drive / Tsinghua Cloud |
| DynamicSwin-B/0.7 | swin-b |
0.7 | 83.43 | 96.45 | 12.1G | Google Drive / Tsinghua Cloud |
- torch>=1.8.0
- torchvision>=0.9.0
- timm==0.3.2
- tensorboardX
- six
- fvcore
Data preparation: download and extract ImageNet images from http://image-net.org/. The directory structure should be
│ILSVRC2012/
├──train/
│ ├── n01440764
│ │ ├── n01440764_10026.JPEG
│ │ ├── n01440764_10027.JPEG
│ │ ├── ......
│ ├── ......
├──val/
│ ├── n01440764
│ │ ├── ILSVRC2012_val_00000293.JPEG
│ │ ├── ILSVRC2012_val_00002138.JPEG
│ │ ├── ......
│ ├── ......
Model preparation: download pre-trained models if necessary:
| model | url | model | url |
|---|---|---|---|
| DeiT-Small | link | LVViT-S | link |
| DeiT-Base | link | LVViT-M | link |
| ConvNeXt-T | link | Swin-T | link |
| ConvNeXt-S | link | Swin-S | link |
| ConvNeXt-B | link | Swin-B | link |
You can try DynamicViT on Colab 
We also provide a Jupyter notebook where you can run the visualization of DynamicViT.
To run the demo, you need to install matplotlib.
To evaluate a pre-trained DynamicViT model on the ImageNet validation set with a single GPU, run:
python infer.py --data_path /path/to/ILSVRC2012/ --model model_name \
--model_path /path/to/model --base_rate 0.7
To train Dynamic Spatial Sparsification models on ImageNet, run:
(You can train models with different keeping ratio by adjusting 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 can also achieve comparable performance with only 15 epochs training (around 0.1% lower accuracy compared to 30 epochs).
ConvNeXt-T
Train on 8 GPUs:
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
Train on 4 8-GPU nodes:
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
Train on 8 GPUs:
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
Train on 4 8-GPU nodes:
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
Train on 8 GPUs:
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
Train on 4 8-GPU nodes:
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
Train on 8 GPUs:
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
Train on 4 8-GPU nodes:
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
Train on 8 GPUs:
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
Train on 4 8-GPU nodes:
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
Train on 8 GPUs:
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
Train on 4 8-GPU nodes:
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 License
Our code is based on pytorch-image-models, DeiT, LV-ViT, ConvNeXt and Swin-Transformer.
If you find our work useful in your research, please consider citing:
@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}
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent