Requirement

torch>=1.7.0; torchvision>=0.8.0; pyyaml; timm==0.6.13;

dataset: Tiny ImageNet

Pretrain Models

Checkpoints of LightViT-XS and LightViT-S have been released.

Train

torchrun --nproc_per_node=4 train.py -c configs/config.yaml --model lightvit --experiment LightViT

Benchmark

python3 benchmark_onnx.py --model lightvit --input-size 3 224 224 --benchmark_cpu

Acknowledgement

Our code base is partly built with SwiftFormer, EfficientMod, image_classification_sota and EfficientFormer.

Thanks for the great implementations!