TinyViT Finetuning and Inference

TinyViT is a new family of tiny and efficient vision transformers pretrained on large-scale datasets with our proposed fast distillation framework. The central idea is to transfer knowledge from large pretrained models to small ones. The logits of large teacher models are sparsified and stored in disk in advance to save the memory cost and computation overheads.

🚀 TinyViT with only 21M parameters achieves 84.8% top-1 accuracy on ImageNet-1k, and 86.5% accuracy under 512x512 resolutions.

this repo uses the pretrained checkpoints of the original authors but provides a more compact solution for finetuning and inference.

Model Zoo

download pth file and put it into ./checkpoints/

Model	Pretrain	Input	Acc@1	Acc@5	#Params	MACs	FPS	22k Model	1k Model
TinyViT-5M	IN-22k	224x224	80.7	95.6	5.4M	1.3G	3,060	link/config/log	link/config/log
TinyViT-11M	IN-22k	224x224	83.2	96.5	11M	2.0G	2,468	link/config/log	link/config/log
TinyViT-21M	IN-22k	224x224	84.8	97.3	21M	4.3G	1,571	link/config/log	link/config/log
TinyViT-21M-384	IN-22k	384x384	86.2	97.8	21M	13.8G	394	-	link/config/log
TinyViT-21M-512	IN-22k	512x512	86.5	97.9	21M	27.0G	167	-	link/config/log
TinyViT-5M	IN-1k	224x224	79.1	94.8	5.4M	1.3G	3,060	-	link/config/log
TinyViT-11M	IN-1k	224x224	81.5	95.8	11M	2.0G	2,468	-	link/config/log
TinyViT-21M	IN-1k	224x224	83.1	96.5	21M	4.3G	1,571	-	link/config/log

TinyViT Fine-tuning Script

Features:

• Load pretrained weights from checkpoint
• Two-stage training: frozen backbone (just head) + full fine-tuning (deep training)
• Early stopping
• TensorBoard logging and Confusion matrix
• Mixed precision training

improvements based on the official TinyViT implementation:

1. Learning Rate Strategy (Major Improvement)

• Cosine annealing with warmup
• Layer-wise learning rate decay (0.8)
• warmup periods (1 epoch for stage 1, 2 epochs for stage 2)
• Low base learning rates (1e-3 for head, 2.5e-4 for deep training)

2. Weight Decay Optimization

• Selective weight decay - excludes bias and normalization layers
• low weight decay (1e-8) for fine-tuning vs standard 0.05
• Proper parameter grouping following TinyViT's approach

3. Training Stability Improvements

• Gradient clipping (max_norm=5.0) prevents exploding gradients
• BatchNorm in eval mode during training
• Better mixed precision handling with unscaling for gradient clipping

4. Enhanced Data Augmentation

• A.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.8)
• A.Rotate(limit=15, p=0.5)
• A.CoarseDropout(max_holes=8, max_height=32, max_width=32, p=0.25) # Random erasing
• A.OneOf([GaussNoise, GaussianBlur], p=0.2) # noise augmentation

5. Training Configuration Updates

• 15 epochs frozen + 5 epochs fine-tuning
• Increased patience (7)

6. Key Functions

1. set_weight_decay() - Proper weight decay handling
2. get_cosine_scheduler_with_warmup() - Advanced LR scheduling
3. Enhanced train_epoch() with gradient clipping and BN eval mode

Citation

If this repo is helpful for you, please consider to cite it. 📣 Thank you! :)

@InProceedings{tiny_vit,
  title={TinyViT: Fast Pretraining Distillation for Small Vision Transformers},
  author={Wu, Kan and Zhang, Jinnian and Peng, Houwen and Liu, Mengchen and Xiao, Bin and Fu, Jianlong and Yuan, Lu},
  booktitle={European conference on computer vision (ECCV)},
  year={2022}
}

License

• License