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.

Arch	Baseline	25%		30%		40%		50%
		1.7.1	1.12.1	1.7.1	1.12.1	1.7.1	1.12.1	1.7.1	1.12.1
Base	81.8	-	-	-	-	81.83	81.77	81.33	81.46
Small	79.9	80.31	80.33	79.90	79.89	-	-	-	-
Tiny	72.2	72.94	72.79	71.90	71.88	-	-	-	-

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 .

Arch	Remove Ratio	w/o ToMe	Started idx	r	w ToMe
Base	40%	81.77	1	24	81.58
			1	28	81.42
	50%	81.46	0	14	81.28
Small	25%	80.33	1	22	79.86
	30%	79.89	1	19	79.62
Tiny	25%	72.79	1	19	72.35
	30%	71.88	1	14	71.7

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.