---
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.