How to launch a training of CAIT models ?

[elias-ramzi]

Hi, thank you for the great repo!

I was wondering what were the commands to launch the training to reproduce the results of the CAIT models.
For now I am mostely interested in the XS and XXS models with 224 resolution.

Is changing drop path and layer scale initialization sufficient?

Also dropout should be at zero?

for now I have the following, after adding the following import: import cait_models in main.py:

python main.py \
--data-path Imagenet \
--model cait_XXS24_224 \
--batch-size 1024 \
--lr 0.001 \
--warmup-lr 0.00001 \
--drop-path 0.05

Thank you!
Elias