Vision Transformer for CIFAR-10
A Vision Transformer (ViT) model trained from scratch on the CIFAR-10 dataset, achieving 82.08% test accuracy.
Model Description
This model implements the Vision Transformer architecture, which processes images as sequences of patches rather than using convolutional layers. The model splits input images into fixed-size patches, linearly embeds them, adds positional encodings, and processes them through multiple transformer encoder layers.
Architecture Details:
- Image Size: 32x32 (CIFAR-10)
- Patch Size: 4x4
- Number of Patches: 64
- Embedding Dimension: 192
- Attention Heads: 3
- Transformer Layers: 12
- MLP Hidden Size: 768
- Total Parameters: 5.36M
Performance
Overall Metrics:
- Test Accuracy: 82.08%
- Test Loss: 1.0026
Per-Class Accuracy:
| Class | Accuracy |
|---|
| Airplane | 85.50% |
| Automobile | 92.80% |
| Bird | 77.70% |
| Cat | 65.90% |
| Deer | 78.40% |
| Dog | 73.10% |
| Frog | 85.50% |
| Horse | 86.60% |
| Ship | 88.80% |
| Truck | 86.50% |
Intended Use
This model is designed for image classification on CIFAR-10 or similar low-resolution datasets. It can be used for:
- Educational purposes to understand Vision Transformer architecture
- Research on transformer-based vision models
- Transfer learning for similar image classification tasks
- Benchmarking and comparison studies
Training Details
Dataset: CIFAR-10
- Training samples: 50,000 images
- Test samples: 10,000 images
- Image dimensions: 32x32x3
- Classes: 10 (airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck)
Training Configuration:
- Optimizer: AdamW
- Learning rate schedule: Cosine annealing with warmup
- Mixed precision training: Enabled
- Label smoothing: 0.1
- Data augmentation: Standard CIFAR-10 augmentations
How to Use
Load the checkpoint using PyTorch:
1import torch
2from vision_transformers.src.models import VisionTransformer
3from vision_transformers.src.config import ViTConfig
4
5# Initialize model
6config = ViTConfig()
7model = VisionTransformer(config)
8
9# Load checkpoint
10checkpoint = torch.load('best_model.pth')
11model.load_state_dict(checkpoint['model_state_dict'])
12model.eval()
13
14# Inference
15with torch.no_grad():
16 logits, _ = model(images)
17 predictions = torch.argmax(logits, dim=1)
Input: Tensor of shape (batch_size, 3, 32, 32) with values normalized to [0, 1]
Output: Logits of shape (batch_size, 10) for CIFAR-10 classes
Limitations
- Trained specifically on CIFAR-10 low-resolution images (32x32)
- Performance varies significantly by class (65.90% for cats vs 92.80% for automobiles)
- Not suitable for high-resolution images without architecture modifications
- May not generalize well to out-of-distribution images
Technical Specifications
Model Architecture:
- Patch Embedding: Conv2d projection with learnable positional embeddings
- Transformer Encoder: Pre-norm architecture with multi-head self-attention
- Classification Head: Layer normalization followed by linear projection
- Activation: GELU in MLP blocks
- Dropout: 0.1 in both attention and MLP layers
Checkpoint Contents:
The .pth file contains:
model_state_dict: Model weightsepoch: Training epoch numberoptimizer_state_dict: Optimizer stateloss: Training lossaccuracy: Validation accuracy
Citation
If you use this model, please cite the original Vision Transformer paper:
1@article{dosovitskiy2020vit,
2 title={An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale},
3 author={Dosovitskiy, Alexey and Beyer, Lucas and Kolesnikov, Alexander and Weissenborn, Dirk and Zhai, Xiaohua and Unterthiner, Thomas and Dehghani, Mostafa and Minderer, Matthias and Heigold, Georg and Gelly, Sylvain and others},
4 journal={ICLR},
5 year={2021}
6}
License
MIT License
