ViT Image Classifier
Vision Transformer From Scratch
An image classifier built from scratch to understand how Vision Transformers work
- 3 patch-embedding strategies
Built with
PyTorch
PoetryViTtorchvision
scikit-learn
StreamlitMatplotlib
Overview
Vision Transformers classify images by cutting them into patches and treating those patches like words in a sentence. I wanted to know what each part of that design actually contributes, so I built a Vision Transformer from scratch in PyTorch rather than importing one.
I wrote each component myself: three patch-embedding strategies, learnable position encodings, multi-head self-attention, a custom GELU activation, and a BERT-style [CLS] token whose final state feeds the classifier. I developed it with the Moroccan Data Scientists community's plant-disease team and applied it to tomato and potato leaf disease images.
On a held-out test split covering four tomato leaf classes, the model reached about 81% accuracy (the repo README reports 81.6%). It is an educational-scale model, not benchmarked against large pretrained ViTs.
The question
What does each part of a Vision Transformer actually contribute — can I build patch embedding, the class token, attention, and fine-tuning from scratch and still classify well?
Key features
- Offers three patch-embedding strategies: a Conv2D projection, LayerNorm on flattened patches before projection, and a pure tensor-reshape approach.
- Prepends a learnable [CLS] token and classifies from its final representation, as in BERT.
- Implements its own GELU activation and LayerNorm with optional bias.
- Uses heavy augmentation for small data: random crops, flips, rotations, color jitter, AutoAugment, and random erasing.
- Trains with AdamW, cosine learning-rate decay, label smoothing, and stratified train/validation/test splits.
- Includes a Streamlit app that loads a checkpoint and classifies an uploaded leaf image.
How it works
Patch the image
Each image is split into 16x16 patches and projected into embedding vectors, with a learnable [CLS] token added in front.
Add positions
Learnable position encodings tell the model where each patch came from, since attention alone ignores order.
Encode
A stack of Transformer blocks applies multi-head self-attention and feed-forward layers so every patch can attend to every other.
Classify
After a final LayerNorm, the [CLS] token's vector goes through a linear head that outputs class scores.
Train and fine-tune
Separate train and fine-tune entry points support a transfer-learning workflow, with live loss and accuracy plots during training.
Examples
Tech stack
Every model component and the training loop
Dependency and environment management
ViT
The architecture rebuilt from scratch
Augmentation pipeline, including AutoAugment
Stratified train/validation/test splits
Upload-and-classify demo app
Live training curves and prediction grids
What I learned
- The patch-tokenization choice changes everything downstream
- The [CLS] representation paradigm transfers cleanly from BERT to vision
Limitations
- Educational scale; not benchmarked against large pretrained ViTs