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.

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%
Arch	Baseline	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
	40%	81.77	1	28	81.42
	50%	81.46	0	14	81.28
Small	25%	80.33	1	22	79.86
Small	30%	79.89	1	19	79.62
Tiny	25%	72.79	1	19	72.35
Tiny	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.
[x] 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.

sihaoevery / lambda_vit

readme