vision-transformers-cifar10

This is your go-to playground for training Vision Transformers (ViT) and its related models on CIFAR-10, a common benchmark dataset in computer vision.

The whole codebase is implemented in Pytorch, which makes it easier for you to tweak and experiment. Over the months, we've made several notable updates including adding different models like ConvMixer, CaiT, ViT-small, SwinTransformers, and MLP mixer. We've also adapted the default training settings for ViT to fit better with the CIFAR-10 dataset.

Using the repository is straightforward - all you need to do is run the train_cifar10.py script with different arguments, depending on the model and training parameters you'd like to use.

Thanks, the repo has been used in 10+ papers!

Please use this citation format if you use this in your research.

@misc{yoshioka2024visiontransformers,
  author       = {Kentaro Yoshioka},
  title        = {vision-transformers-cifar10: Training Vision Transformers (ViT) and related models on CIFAR-10},
  year         = {2024},
  publisher    = {GitHub},
  howpublished = {\url{https://github.com/kentaroy47/vision-transformers-cifar10}}
}

Updates

• Added ConvMixer implementation. Really simple! (2021/10)

• Added wandb train log to reproduce results. (2022/3)

• Added CaiT and ViT-small. (2022/3)

• Added SwinTransformers. (2022/3)

• Added MLP mixer. (2022/6)

• Changed default training settings for ViT.

• Fixed some bugs and training settings (2024/2)

• Added onnx and torchscript model exports. (2024/12)

Usage example

python train_cifar10.py # vit-patchsize-4

python train_cifar10.py --size 48 # vit-patchsize-4-imsize-48

python train_cifar10.py --patch 2 # vit-patchsize-2

python train_cifar10.py --net vit_small --n_epochs 400 # vit-small

python train_cifar10.py --net vit_timm # train with pretrained vit

python train_cifar10.py --net convmixer --n_epochs 400 # train with convmixer

python train_cifar10.py --net mlpmixer --n_epochs 500 --lr 1e-3

python train_cifar10.py --net cait --n_epochs 200 # train with cait

python train_cifar10.py --net swin --n_epochs 400 # train with SwinTransformers

python train_cifar10.py --net res18 # resnet18+randaug

Results..

	Accuracy	Train Log
ViT patch=2	80%
ViT patch=4 Epoch@200	80%	Log
ViT patch=4 Epoch@500	88%	Log
ViT patch=8	30%
ViT small	80%
MLP mixer	88%
CaiT	80%
Swin-t	90%
ViT small (timm transfer)	97.5%
ViT base (timm transfer)	98.5%
ConvMixerTiny(no pretrain)	96.3%	Log
resnet18	93%
resnet18+randaug	95%	Log

Used in..

• Vision Transformer Pruning arxiv github
• Understanding why ViT trains badly on small datasets: an intuitive perspective arxiv
• Training deep neural networks with adaptive momentum inspired by the quadratic optimization arxiv
• Moderate coreset: A universal method of data selection for real-world data-efficient deep learning

Model Export

This repository supports exporting trained models to ONNX and TorchScript formats for deployment purposes. You can export your trained models using the export_models.py script.

Basic Usage

python export_models.py --checkpoint path/to/checkpoint --model_type vit --output_dir exported_models