liangnanying/vit-scratch-cifar10

ViT from Scratch — CIFAR-10 (IT3103 Week 15)

A small Vision Transformer trained from scratch on CIFAR-10 for teaching purposes. Final test accuracy: 77.13%.

This is a plain PyTorch nn.Module (see the ViT/TransformerBlock/PatchEmbedding classes in the course notebook), not a 🤗 Transformers model. To use it, recreate the architecture from the notebook and load the weights:

import json, torch
from huggingface_hub import hf_hub_download

cfg = json.load(open(hf_hub_download("liangnanying/vit-scratch-cifar10", "config.json")))
model = ViT(  # <-- the ViT class from the notebook
    img_size=cfg["img_size"], patch_size=cfg["patch_size"],
    num_classes=cfg["num_classes"], embed_dim=cfg["embed_dim"],
    depth=cfg["depth"], num_heads=cfg["num_heads"], dropout=cfg["dropout"],
)
sd = torch.load(hf_hub_download("liangnanying/vit-scratch-cifar10", "vit_scratch_cifar10.pt"), map_location="cpu")
model.load_state_dict(sd)
model.eval()

Config: {"architecture": "ViT (from scratch, custom nn.Module)", "img_size": 32, "patch_size": 4, "in_channels": 3, "num_classes": 10, "embed_dim": 256, "depth": 4, "num_heads": 8, "mlp_ratio": 4.0, "dropout": 0.1, "classes": ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"], "normalize_mean": [0.4914, 0.4822, 0.4465], "normalize_std": [0.247, 0.2435, 0.2616], "final_test_acc": 77.13}