ViT CIFAR-10 Classifier
Model Description
This model is a Vision Transformer (ViT) fine-tuned for image classification on the CIFAR-10 dataset.
It recognizes 10 image categories:
- airplane
- automobile
- bird
- cat
- deer
- dog
- frog
- horse
- ship
- truck
Base Model
google/vit-base-patch16-224-in21k
Dataset
CIFAR-10 (uoft-cs/cifar10).
The original dataset contains 50,000 training images and 10,000 test images.
For this experiment, we used:
- Training samples: 2,000
- Validation samples: 500
Training Configuration
| Parameter | Value |
|---|---|
| Epochs | 1 |
| Batch size | 8 |
| Learning rate | 2e-5 |
| Weight decay | 0.01 |
| Image size | 224 × 224 |
| Hardware | NVIDIA T4 GPU |
Evaluation Results
Evaluation on 500 validation images:
| Metric | Value |
|---|---|
| Accuracy | 90.0% |
| Validation loss | 1.5738 |
Per-class F1 Scores
| Class | F1-score |
|---|---|
| airplane | 0.986 |
| automobile | 0.933 |
| bird | 0.750 |
| cat | 0.809 |
| deer | 0.931 |
| dog | 0.828 |
| frog | 0.903 |
| horse | 0.941 |
| ship | 0.990 |
| truck | 0.941 |
Limitations
- The model was fine-tuned on only 2,000 training images.
- Evaluation was performed on a 500-image validation subset.
- Performance varies between classes, especially for birds.
- The reported results are not a full test-set evaluation.
Experiment Tracking
Training parameters, metrics, confusion matrix, and per-class metrics were recorded using MLflow.
Intended Use
Educational image classification experiments and demonstrations using CIFAR-10.