ViT — Vision Transformer Framework

A lightweight, PyTorch-based framework for training Vision Transformers (ViT) on your own image classification datasets — with automatic data splitting, checkpointing, logging, and evaluation built in.

Published on PyPI as vit-hk.

Installation

pip install vit-hk

This installs the package along with its dependencies: torch, torchvision, scikit-learn, and tqdm.

Note: the PyPI package name is vit-hk, but you import it as ViT in Python (see below).

What's in this repo

ViT/
├── src/ViT/
│   ├── core.py       # ViT model architecture (PatchEmbedding, MultiHeadSelfAttention, MLP, TransformerBlock, ViT)
│   ├── data.py        # Dataset loading + automatic train/val/test splitting
│   ├── trainer.py      # Training loop, checkpointing, evaluation
│   └── utils.py        # Logging setup + metrics (accuracy, precision, recall, F1, confusion matrix)
├── examples/
│   ├── train.py         # End-to-end training script (CLI)
│   └── smoke_test.py     # Generates a tiny fake dataset to sanity-check the pipeline
├── tests/
│   └── test_core.py       # Unit tests
├── pyproject.toml
└── README.md

Quick start: use the model directly

import torch
from ViT import ViT

model = ViT(
    img_size=224,
    patch_size=16,
    num_classes=1000,
    embed_dim=768,
    depth=12,
    num_heads=12,
)

x = torch.randn(1, 3, 224, 224)  # (batch, channels, height, width)
logits = model(x)                # -> (1, 1000)

You can also import individual building blocks to construct your own custom transformer variants:

from ViT import TransformerBlock, PatchEmbedding, MultiHeadSelfAttention, MLP

Full pipeline: train on your own dataset

1. Organize your data

Put your images in one folder, with one subfolder per class:

data/
├── cat/
│   ├── img1.jpg
│   └── img2.jpg
├── dog/
│   ├── img1.jpg
│   └── img2.jpg

2. Run training

python examples/train.py --data_dir path/to/data --epochs 20 --save_every 5

This will automatically:

• Split your data into train / validation / test sets (default 70/15/15, configurable)
• Train the model, printing loss and validation metrics every epoch
• Save a checkpoint every --save_every epochs, plus always keep the best model (highest validation accuracy) and the final model
• Log everything to both the console and a training.log file
• Run a full evaluation on the held-out test set at the end (accuracy, precision, recall, F1, confusion matrix)

CLI options

3. Use it programmatically instead of the CLI

from ViT import ViT, Trainer, prepare_dataloaders

train_loader, val_loader, test_loader, class_names = prepare_dataloaders(
    data_dir="path/to/data",
    img_size=224,
    batch_size=32,
)

model = ViT(img_size=224, num_classes=len(class_names), embed_dim=384, depth=6, num_heads=6)

trainer = Trainer(model, checkpoint_dir="checkpoints", save_every=5, class_names=class_names)
trainer.fit(train_loader, val_loader, epochs=20)
trainer.evaluate(test_loader, split_name="test")

4. Loading a saved checkpoint later

trainer.load_checkpoint("checkpoints/best.pt")

Sanity-checking before training on real data

Generate a tiny synthetic dataset to confirm the whole pipeline runs on your machine:

python examples/smoke_test.py
python examples/train.py --data_dir fake_data --epochs 2 --save_every 1

Development / running tests

pip install -e .
pytest tests/ -v

License

MIT — see pyproject.toml for author/license details.