Awesome
GreenMIM
This is the official PyTorch implementation of the NeurIPS 2022 paper Green Hierarchical Vision Transformer for Masked Image Modeling. GreenMIM consists of two key desgins, Group Window Attention
and Sparse Convolution
. It offers 2.7x faster pre-training and competitive performance on hierarchical vision transformers, e.g., Swin/Twins Transformers.
Citation
If you find our work interesting or use our code/models, please cite:
@article{huang2022green,
title={Green Hierarchical Vision Transformer for Masked Image Modeling},
author={Huang, Lang and You, Shan and Zheng, Mingkai and Wang, Fei and Qian, Chen and Yamasaki, Toshihiko},
journal={Thirty-Sixth Conference on Neural Information Processing Systems},
year={2022}
}
News
- 2023.01: We have refactor the structure of this codebase, supporting most, if not any, vision transformer backbones with various input resolutions. Checkout our implementation of GreenMIM with Twins Transformer here.
Catalogs
- Pre-trained checkpoints
- Pre-training code for
Swin Transformer
andTwins Transformer
- Fine-tuning code
Pre-trained Models
<table><tbody> <!-- START TABLE --> <!-- TABLE HEADER --> <th valign="bottom"></th> <th valign="bottom">Swin-Base (Window 7x7)</th> <th valign="bottom">Swin-Base (Window 14x14)</th> <th valign="bottom">Swin-Large (Window 14x14)</th> <!-- TABLE BODY --> <tr><td align="left">pre-trained checkpoint</td> <td align="center"><a href="https://drive.google.com/file/d/1vCt7QN3rNC7hmWlWYomqfhjUqN-PvR7a/view?usp=sharing">Download</a></td> <td align="center"><a href="https://drive.google.com/file/d/1P1dAdcZtSEGWFQy5GeeJdfGTqesSAES9/view?usp=sharing">Download</a></td> <td align="center"><a href="https://drive.google.com/file/d/1Tw1KeGviVWxbVt3h1TT7BxTX1aLeT-Nm/view?usp=sharing">Download</a></td> </tbody></table>Pre-training
The pre-training scripts are given in the scripts/
folder. The scripts with names start with 'run*' are for non-slurm users while the others are for slurm users.
For Non-Slurm Users
To train a Swin-B with on a single node with 8 GPUs.
PORT=23456 NPROC=8 bash scripts/run_greenmim_swin_base.sh
For Slurm Users
To train a Swin-B with on a single node with 8 GPUs.
bash scripts/srun_greenmim_swin_base.sh [Partition] [NUM_GPUS]
Fine-tuning on ImageNet-1K
Model | #Params | Pre-train Resolution | Fine-tune Resolution | Config | Acc@1 (%) |
---|---|---|---|---|---|
Swin-B (Window 7x7) | 88M | 224x224 | 224x224 | Config | 83.8 |
Swin-L (Window 14x14) | 197M | 224x224 | 224x224 | Config | 85.1 |
Currently, we directly use the code of SimMIM for fine-tuning, please follow their instructions to use the configs. NOTE that, due to the limited computing resource, we use a batch size of a batch size of 768 (48 x 16) for fine-tuning.
Acknowledgement
This code is based on the implementations of MAE, SimMIM, BEiT, SwinTransformer, Twins Transformer, and DeiT.
License
This project is under the CC-BY-NC 4.0 license. See LICENSE for details.