Awesome

MLP Can Be A Good Transformer Learner (CVPR2024)

Overview

• The implementation of Base architecture (deit_base_patch16_224_attn) can be found in models_v4.py. For Small (deit_small_patch16_224_copy_lambda) and Tiny (deit_tiny_patch16_224_copy_lambda) architectures, they are in models_small_timm3_lambda.py.
• We copy the code from timm 0.3.2 to implement the baseline for small and tiny archs in models_small_timm3.py, which further results in our model for small and tiny archs in models_small_timm3_lambda.py. However, the model deit_base_patch16_224_copy_lambda has not been tested.

Note

• There is a BUG when saving the best models in main_lambda.py because it will save the model with lambda>0. It's easy to fix this bug if you want to save the best checkpoint with lambda=0.
• Additionally, we save the checkpoints at 280, 290 and 295 epochs (lambda>0). If you don't need them, just comment the code snippets.
• Hence, we use the final checkpoint for evaluation. Again, if you want to use the best checkpoint with lambda=0, it's easy to fix the bug on your own.

Requirements

• timm==0.4.12
• einops
• torchprofile
• fvcore

Note that we use torch==1.7.1 for training. To incorparate with ToMe, we use torch==1.12.1.

Checkpoints

We provide some checkpoints for reference. Here the prefix indicates the architectures while the suffix indicates which attention layers are removed.

• base_01346.pth
• base_013469.pth
• small_024.pth
• small_0246.pth
• tiny_024.pth
• tiny_0246.pth

Performance

We found that the same code and checkpoint would produce different inference results using different pytorch versions. We still cannot figure out and welcome discussions.

We deploy the ToMe over the normal blocks (indexed by 0, 1, 2, ...). Typically, we use this technique on the normal block started by index 1 and its subsequent normal blocks. The model is evaluated with torch==1.12.1 .

Before getting started

• Please download the DeiT checkpoints.
• Modify the ImageNet path in shell scripts, e.g. script/shrink_base.sh.

Training

We use 8 GPUs with 256 images per GPU.

E.g.

./script/shrink_base.sh

Testing

./script/test.sh

Speed, Params & FLOPs

Please refer to benchmark.py and run

python benchmark.py

To-Do

• Code for segmentation.
• Upload checkpoints.

Issues / Contact

Feel free to create an issue if you get a question or just drop me emails ( sihao.lin@student.rmit.edu.au ).

Acknowledgement

This work is built upon DeiT. Thanks to their awesome work.