A compact Vision Transformer trained from scratch on CIFAR-10, with a full loop from training to a small interactive demo.
Model#
Classic ViT-style stack on 32×32 RGB:
- Patch embed: conv with
patch_size=4→ 8×8 = 64 tokens - multi-head self-attention + MLP blocks (GELU)
- CLS-style classification head for 10 CIFAR classes
Demo checkpoint config (in app.py): embed_dim=256, depth=8, n_heads=8, light dropout.
Pipeline#
| Script | Role |
|---|---|
train.py | CIFAR-10 train/val, RandAugment, AdamW, cosine LR, grad clip, best ckpt |
evaluate.py | val loss and accuracy from saved weights |
inference.py | single-image predict (resize + CIFAR normalize) |
app.py | Gradio UI with class confidence bars |
push_to_hub.py | upload model and card via huggingface_hub |
launch_sagemaker.py | cloud training helper |
Training defaults are a practical small-ViT recipe: batch 128, ~50 epochs, LR 3e-4, weight decay 0.05, RandomCrop / flip / RandAugment. Channel layout also works with SageMaker env vars.
Demo#
Gradio shows softmax confidences over:
plane, car, bird, cat, deer, dog, frog, horse, ship, truck
Example images live under images/. Checkpoint path defaults to model/best_model.pth.
Why it exists#
Same idea as the Qwen lab: implement the architecture, own the train loop, ship something people can click. CIFAR-10 keeps the bill low while still stressing patch embedding, attention depth, and regularization on small images.