File size: 5,224 Bytes
07d92db 205ea67 20f20f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | ---
license: mit
datasets:
- ILSVRC/imagenet-1k
language:
- en
metrics:
- accuracy
pipeline_tag: image-classification
library_name: transformers
tags:
- vit
- nano
- patch16
- img224
---
# CustomViT-Nano: 4.24M Parameter Compact Vision Transformer
**CustomViT-Nano** is a compact, modernized Vision Transformer architecture designed for efficient ImageNet-1K image classification under a small parameter budget. With only **4.24M trainable parameters**, the model combines a lightweight convolutional stem with modern Transformer components including **2D Rotary Positional Embeddings**, **Pre-RMSNorm**, **SwiGLU feed-forward layers**, and **PyTorch SDPA attention**. The model is designed to deliver strong classification performance while remaining significantly smaller than standard large Vision Transformer baselines.
---
## Key Architectural Features
CustomViT-Nano modernizes a small Vision Transformer design using several efficiency-focused architectural components.
| Component | Design in CustomViT-Nano |
|---|---|
| **Patch Embedding** | Multi-stage convolutional stem instead of single large patchify projection |
| **Stem Activation** | GELU |
| **Token Layout** | 14 × 14 patch tokens + CLS token |
| **Normalization** | Pre-RMSNorm inside Transformer blocks |
| **Attention** | Multi-head attention using PyTorch scaled dot-product attention |
| **Position Encoding** | 2D Rotary Positional Embeddings for image patch grids |
| **MLP Block** | SwiGLU gated feed-forward network |
| **Classifier** | CLS-token classification head |
---
## ConvStem Design
Instead of directly projecting `16 × 16` image patches with one large-stride convolution, CustomViT-Nano uses a progressive convolutional stem:
```text
224 × 224 × 3
↓
112 × 112 × 32
↓
56 × 56 × 64
↓
28 × 28 × 128
↓
14 × 14 × 224
```
This gives the model a stronger local visual inductive bias before global Transformer reasoning.
---
## Benchmark & Evaluation
- **Evaluation Dataset**: ImageNet-1K validation set
- **Total Evaluation Images**: 50,000
- **Input Resolution**: 224 × 224
- **Number of Classes**: 1000
| Model | Parameters | Top-1 Accuracy | Top-5 Accuracy |
|---|---:|---:|---:|
| **CustomViT-Nano** | **4.24M** | **63.60%** | **84.93%** |
| **Google ViT-B/16** | **86.6M** | **80.31%** | **95.49%** |
---
## Parameter Efficiency Comparison
CustomViT-Nano is approximately **20.4× smaller** than the reference Google ViT-B/16 model.
```text
Google ViT-B/16: 86.6M parameters
CustomViT-Nano: 4.24M parameters
```
Parameter reduction:
```text
86.6M / 4.24M ≈ 20.4× smaller
```
Despite using only around **4.9%** of the parameters of the 86.6M ViT baseline, CustomViT-Nano achieves:
- **63.60% Top-1 Accuracy**
- **84.93% Top-5 Accuracy**
on the ImageNet-1K validation set.
---
## Target Use Cases & Applications
CustomViT-Nano is suitable for scenarios where a compact visual classifier is preferred over a large transformer model.
1. **Compact Image Classification**
Lightweight ImageNet-style classification with a transformer-based architecture.
2. **Edge & Resource-Constrained Vision**
Useful for environments where model size and memory footprint are important constraints.
3. **Educational Vision Transformer Research**
A compact architecture for studying ConvStem patch embeddings, 2D RoPE, SDPA attention, RMSNorm, and SwiGLU inside a full ImageNet-scale classifier.
4. **Backbone Experiments**
Can be used as a small image encoder backbone for downstream classification or transfer-learning experiments.
5. **Efficient Model Baselines**
Useful as a compact baseline for experiments involving distillation, pruning, quantization, or architecture search.
---
## How to Use
### Fast Inference with Hugging Face `pipeline`
```python
from transformers import pipeline
classifier = pipeline(
"image-classification",
model="kd13/vit-nano-patch16-224",
trust_remote_code=True
)
results = classifier(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/hub/parrots.png"
)
for pred in results:
print(f"Label: {pred['label']} | Score: {pred['score']:.4f}")
```
---
## Inference with PIL Image
```python
from PIL import Image
from transformers import pipeline
classifier = pipeline(
"image-classification",
model="kd13/vit-nano-patch16-224",
trust_remote_code=True
)
image = Image.open("image.jpg").convert("RGB")
results = classifier(image)
for pred in results:
print(f"Label: {pred['label']} | Score: {pred['score']:.4f}")
```
---
## Limitations
- The model is smaller than standard ViT-B models and therefore has lower absolute ImageNet accuracy.
- It is optimized for image classification, not detection, segmentation, captioning, or multimodal tasks.
- Performance may vary on images that differ significantly from ImageNet-style natural images.
- For maximum accuracy, larger models or teacher-distilled variants may perform better.
---
## Disclaimer
This model is intended for research, experimentation, and efficient image-classification use cases. It should be evaluated carefully before use in production or safety-critical applications. |