Training and evaluation code for LIT-S, LIT-M and LIT-B.

First, activate your python environment

conda activate lit

Make sure you have the correct ImageNet DATA_PATH in config/*.yaml.

To train LIT-S:

bash scripts/lit-s.sh [GPUs] 

To train LIT-M:

bash scripts/lit-m.sh [GPUs] 

To train LIT-B:

bash scripts/lit-b.sh [GPUs] 

Note: We use a total batch size of 1024 for all experiments on ImageNet. Therefore, you may want to use a different batch size by editing BATCH_SIZE in configs/*.yaml. For example, by setting BATCH_SIZE to 64 and training with 8 GPUs, your total batch size is 512.

We provide scripts to evaluate LIT-S, LIT-M and LIT-B. To evaluate a model, you can run

bash scripts/lit-b-eval.sh [GPUs] [path/to/checkpoint]

For example, to evaluate LIT-B with 1 GPU, you can run:

bash scripts/lit-b-eval.sh 1 checkpoint/lit_b.pth

This should give

* Acc@1 83.366 Acc@5 96.254
Accuracy of the network on the 50000 test images: 83.4%

Result could be slightly different based on you environment.