Metadata-Version: 2.4
Name: vit-hk
Version: 0.1.3
Summary: Implementation of Vision Transformer (ViT) in PyTorch.
Author-email: Hamza Khan <hamxa678@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/hamxa678/ViT
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Dynamic: license-file

# 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

```bash
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

```python
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:

```python
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

```bash
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

| Argument | Default | Description |
|---|---|---|
| `--data_dir` | *(required)* | Path to your dataset folder |
| `--img_size` | 224 | Image resize dimension |
| `--batch_size` | 32 | Training batch size |
| `--epochs` | 20 | Number of training epochs |
| `--save_every` | 5 | Save a checkpoint every N epochs |
| `--lr` | 3e-4 | Learning rate |
| `--train_split` | 0.7 | Fraction of data used for training |
| `--val_split` | 0.15 | Fraction of data used for validation |
| `--test_split` | 0.15 | Fraction of data used for testing |
| `--checkpoint_dir` | `checkpoints` | Where model checkpoints are saved |

### 3. Use it programmatically instead of the CLI

```python
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

```python
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:

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

---

## Development / running tests

```bash
pip install -e .
pytest tests/ -v
```

---

## License

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