Efficient Vision Transformers and CNNs with Dynamic Spatial Sparsification
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)]
🔥 Updates
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.
Image Examples
Video Examples
Model Zoo
We provide our DynamicViT models pretrained on ImageNet:
| name | model | rho | acc@1 | acc@5 | FLOPs | url |
|---|---|---|---|---|---|---|
| DynamicViT-256/0.7 | deit-256 |
0.7 | 76.53 | 93.12 | 1.3G | Google Drive / Tsinghua Cloud |
| DynamicViT-384/0.7 | deit-s |
0.7 | 79.32 | 94.68 | 2.9G | Google Drive / Tsinghua Cloud |
| DynamicViT-LV-S/0.5 | lvvit-s |
0.5 | 81.97 | 95.76 | 3.7G | Google Drive / Tsinghua Cloud |
| DynamicViT-LV-S/0.7 | lvvit-s |
0.7 | 83.08 | 96.25 | 4.6G | Google Drive / Tsinghua Cloud |
| DynamicViT-LV-M/0.7 | lvvit-m |
0.7 | 83.82 | 96.58 | 8.5G | Google Drive / Tsinghua Cloud |
| 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 |
Usage
Requirements
- 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 |
Demo
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.
Evaluation
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
Training
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
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
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
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
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
License
MIT License
Acknowledgements
Our code is based on pytorch-image-models, DeiT, LV-ViT, ConvNeXt and Swin-Transformer.
Citation
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}