ViT Cats, Dogs and Pandas Classifier
Vision Transformer image classifier for three classes:
0: cats1: dogs2: panda
Base model
google/vit-base-patch16-224
The classifier uses the output representation of the first [CLS] token and a
Linear(768, 3) classification head.
Dataset
Melisa13/Animals_dataset
| Split | Images | Per class |
|---|---|---|
| Train | 408 | 136 |
| Validation | 72 | 24 |
| Test | 120 | 40 |
Exact image duplicates between splits: 0.
Image preprocessing
- RGB input
- Resize and crop to
224×224 - Patch size:
16×16 - 196 image patches
- 197 tokens including
[CLS] - Normalization mean:
(0.5, 0.5, 0.5) - Normalization standard deviation:
(0.5, 0.5, 0.5)
Fine-tuning
The following components were trainable:
- Transformer blocks 10 and 11
- Final LayerNorm
Linear(768, 3)classifier
| Configuration | Value |
|---|---|
| Total parameters | 85,800,963 |
| Trainable parameters | 14,179,587 |
| Trainable share | 16.53% |
| Epochs completed | 7 |
| Batch size | 16 |
| Learning rate | 2e-5 |
| Weight decay | 0.01 |
Results
| Split | Accuracy | Macro F1 |
|---|---|---|
| Validation | 98.61% | 98.61% |
| Test | 100.00% | 100.00% |
Test classification report:
| Class | Precision | Recall | F1 | Support |
|---|---|---|---|---|
| Cats | 1.0000 | 1.0000 | 1.0000 | 40 |
| Dogs | 1.0000 | 1.0000 | 1.0000 | 40 |
| Panda | 1.0000 | 1.0000 | 1.0000 | 40 |
End-to-end latency
Latency includes image preprocessing and model inference with batch size 1.
| Device | Mean | Median | P95 | Throughput |
|---|---|---|---|---|
| Tesla T4 | 21.86 ms | 18.37 ms | 39.92 ms | 45.75 images/s |
| Colab CPU | 599.39 ms | 371.03 ms | 1550.67 ms | 1.67 images/s |
Usage
from PIL import Image
import torch
from transformers import (
AutoImageProcessor,
AutoModelForImageClassification,
)
MODEL_NAME = "Bektur756/vit-cats-dogs-pandas"
processor = AutoImageProcessor.from_pretrained(
MODEL_NAME
)
model = AutoModelForImageClassification.from_pretrained(
MODEL_NAME
)
image = Image.open(
"animal.jpg"
).convert(
"RGB"
)
inputs = processor(
images=image,
return_tensors="pt",
)
with torch.inference_mode():
logits = model(
**inputs
).logits
probabilities = torch.softmax(
logits,
dim=-1,
)[0]
predicted_id = int(
probabilities.argmax()
)
print(
model.config.id2label[
predicted_id
]
)
print(
probabilities.tolist()
)
## Limitations
- The test split contains only 120 images.
- The dataset is small and relatively simple.
- Perfect test accuracy does not imply perfect performance on arbitrary images.
- CPU inference is substantially slower than GPU inference.