Implemented Vision Transformer architechture from scratch!

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+inference_samples.png filter=lfs diff=lfs merge=lfs -text
+training_curves.png filter=lfs diff=lfs merge=lfs -text

.virtual_documents/__notebook_source__.ipynb

@@ -0,0 +1,854 @@
+
+
+
+# 1. Install dependencies (uncomment if running in a fresh environment)
+# !pip install -q torch torchvision datasets matplotlib scikit-learn tqdm
+
+
+
+# 2. Imports
+import os
+import json
+import time
+import math
+import random
+import copy
+from dataclasses import dataclass, field, asdict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+import torchvision.transforms as T
+import torchvision.transforms.functional as TF
+
+import matplotlib.pyplot as plt
+from tqdm.auto import tqdm
+
+from safetensors.torch import save_file, load_file
+
+from PIL import Image
+
+print("Torch:", torch.__version__)
+print("CUDA available:", torch.cuda.is_available())
+
+
+
+# 3. Reproducibility
+SEED = 42
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+
+
+
+# 4. Configuration
+@dataclass
+class Config:
+    # Data
+    image_size: int = 224
+    patch_size: int = 16
+    num_classes: int = 151          # 150 ADE20K classes + 1 background/ignore(0)
+    ignore_index: int = 0           # ADE20K label 0 = unlabeled/background
+
+    # Model
+    embed_dim: int = 384            # ViT-Small-ish for a runnable demo (use 768 for ViT-Base)
+    depth: int = 8
+    num_heads: int = 6
+    mlp_ratio: float = 4.0
+    dropout: float = 0.1
+
+    # Training
+    batch_size: int = 8
+    epochs: int = 10                 # increase for real training (e.g. 64-160)
+    lr: float = 3e-4
+    weight_decay: float = 0.05
+    warmup_epochs: int = 1
+    num_workers: int = 2
+    subset_size: int = None         # e.g. 500 to train fast on a subset; None = full split
+
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+
+cfg = Config()
+print(cfg)
+
+
+
+
+
+
+from datasets import load_dataset
+
+ade20k = load_dataset(
+    "parquet",
+    data_files={
+        "train": "https://huggingface.co/datasets/zhoubolei/scene_parse_150/resolve/refs%2Fconvert%2Fparquet/scene_parsing/train/0000.parquet",
+        "validation": "https://huggingface.co/datasets/zhoubolei/scene_parse_150/resolve/refs%2Fconvert%2Fparquet/scene_parsing/validation/0000.parquet",
+    },
+    features=None,
+)
+
+print(ade20k)
+print(ade20k["train"].column_names)
+
+
+# 6. Visualize a raw sample (image + segmentation mask)
+sample = ade20k["train"][0]
+fig, axes = plt.subplots(1, 2, figsize=(10, 5))
+axes[0].imshow(sample["image"])
+axes[0].set_title("Image")
+axes[0].axis("off")
+
+axes[1].imshow(sample["annotation"], cmap="tab20")
+axes[1].set_title("Annotation (class ids)")
+axes[1].axis("off")
+plt.tight_layout()
+plt.show()
+
+
+
+# 7. Dataset wrapper: resize image + mask, normalize image, return tensors
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+
+class ADE20KSegDataset(Dataset):
+    # Wraps a HF `scene_parse_150` split into (image_tensor, mask_tensor) pairs.
+    def __init__(self, hf_split, image_size=224, train=True, subset_size=None):
+        self.data = hf_split
+        self.image_size = image_size
+        self.train = train
+        self.indices = list(range(len(self.data)))
+        if subset_size is not None:
+            self.indices = self.indices[:subset_size]
+
+        self.img_tf = T.Compose([
+            T.Resize((image_size, image_size), interpolation=T.InterpolationMode.BILINEAR),
+            T.ToTensor(),
+            T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+        ])
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __getitem__(self, idx):
+        item = self.data[self.indices[idx]]
+        image = item["image"].convert("RGB")
+        mask = item["annotation"]  # PIL "L" image, pixel values = class id (0..150)
+
+        # Simple augmentation: random horizontal flip (train only)
+        if self.train and random.random() < 0.5:
+            image = TF.hflip(image)
+            mask = TF.hflip(mask)
+
+        img_t = self.img_tf(image)
+
+        mask = mask.resize((self.image_size, self.image_size), Image.NEAREST)
+        mask_t = torch.from_numpy(np.array(mask)).long()
+
+        return img_t, mask_t
+
+
+train_dataset = ADE20KSegDataset(ade20k["train"], image_size=cfg.image_size,
+                                  train=True, subset_size=cfg.subset_size)
+val_dataset = ADE20KSegDataset(ade20k["validation"], image_size=cfg.image_size,
+                                train=False, subset_size=cfg.subset_size)
+
+print("Train size:", len(train_dataset), "| Val size:", len(val_dataset))
+
+
+
+# 8. DataLoaders
+train_loader = DataLoader(
+    train_dataset, batch_size=cfg.batch_size, shuffle=True,
+    num_workers=cfg.num_workers, pin_memory=True, drop_last=True, persistent_workers=True,
+)
+val_loader = DataLoader(
+    val_dataset, batch_size=cfg.batch_size, shuffle=False,
+    num_workers=cfg.num_workers, pin_memory=True, persistent_workers=True,
+)
+
+imgs, masks = next(iter(train_loader))
+print("Batch image shape:", imgs.shape, "| Batch mask shape:", masks.shape)
+
+
+
+
+
+
+# 9. Patch Embedding
+class PatchEmbedding(nn.Module):
+    def __init__(self, image_size=224, patch_size=16, in_channels=3, embed_dim=768):
+        super().__init__()
+        assert image_size % patch_size == 0
+
+        self.grid_size = image_size // patch_size
+        self.num_patches = self.grid_size ** 2
+
+        self.projection = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, x):
+        # x: [B, C, H, W]
+        x = self.projection(x)       # [B, D, H/P, W/P]
+        x = x.flatten(2)             # [B, D, N]
+        x = x.transpose(1, 2)        # [B, N, D]
+        return x
+
+
+
+# 10. Multi-Head Self-Attention
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads, dropout=0.0):
+        super().__init__()
+        assert embed_dim % num_heads == 0
+
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+
+        self.qkv = nn.Linear(embed_dim, 3 * embed_dim)
+        self.proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        B, N, D = x.shape
+
+        qkv = self.qkv(x)
+        qkv = qkv.reshape(B, N, 3, self.num_heads, self.head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        Q, K, V = qkv[0], qkv[1], qkv[2]
+
+        attention = (Q @ K.transpose(-2, -1)) / (self.head_dim ** 0.5)
+        attention = attention.softmax(dim=-1)
+        attention = self.dropout(attention)
+
+        x = attention @ V
+        x = x.transpose(1, 2)
+        x = x.reshape(B, N, D)
+        x = self.proj(x)
+        return x
+
+
+
+# 11. MLP block
+class MLP(nn.Module):
+    def __init__(self, embed_dim, mlp_ratio=4.0, dropout=0.0):
+        super().__init__()
+        hidden_dim = int(embed_dim * mlp_ratio)
+
+        self.fc1 = nn.Linear(embed_dim, hidden_dim)
+        self.act = nn.GELU()
+        self.dropout1 = nn.Dropout(dropout)
+
+        self.fc2 = nn.Linear(hidden_dim, embed_dim)
+        self.dropout2 = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.dropout1(x)
+        x = self.fc2(x)
+        x = self.dropout2(x)
+        return x
+
+
+
+# 12. Transformer Encoder Block (pre-norm, residual)
+class TransformerEncoderBlock(nn.Module):
+    def __init__(self, embed_dim, num_heads, mlp_ratio=4.0, dropout=0.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.attention = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
+
+        self.norm2 = nn.LayerNorm(embed_dim)
+        self.mlp = MLP(embed_dim, mlp_ratio, dropout)
+
+    def forward(self, x):
+        x = x + self.attention(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+
+
+
+
+# 13. ViT backbone (no classification head — returns patch-token features)
+class ViTBackbone(nn.Module):
+    def __init__(
+        self,
+        image_size=224,
+        patch_size=16,
+        in_channels=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        dropout=0.0,
+        use_cls_token=True,
+    ):
+        super().__init__()
+        self.patch_embedding = PatchEmbedding(image_size, patch_size, in_channels, embed_dim)
+        self.grid_size = self.patch_embedding.grid_size
+        num_patches = self.patch_embedding.num_patches
+        self.use_cls_token = use_cls_token
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if use_cls_token else None
+        num_tokens = num_patches + (1 if use_cls_token else 0)
+        self.pos_embedding = nn.Parameter(torch.zeros(1, num_tokens, embed_dim))
+
+        self.dropout = nn.Dropout(dropout)
+
+        self.blocks = nn.ModuleList([
+            TransformerEncoderBlock(embed_dim, num_heads, mlp_ratio, dropout)
+            for _ in range(depth)
+        ])
+        self.norm = nn.LayerNorm(embed_dim)
+
+        self._init_weights()
+
+    def _init_weights(self):
+        if self.cls_token is not None:
+            nn.init.trunc_normal_(self.cls_token, std=0.02)
+        nn.init.trunc_normal_(self.pos_embedding, std=0.02)
+
+    def forward(self, x):
+        x = self.patch_embedding(x)          # [B, N, D]
+        B = x.shape[0]
+
+        if self.use_cls_token:
+            cls_token = self.cls_token.expand(B, -1, -1)
+            x = torch.cat([cls_token, x], dim=1)
+
+        x = x + self.pos_embedding
+        x = self.dropout(x)
+
+        for block in self.blocks:
+            x = block(x)
+
+        x = self.norm(x)
+
+        if self.use_cls_token:
+            patch_tokens = x[:, 1:]          # drop CLS, keep patch tokens
+        else:
+            patch_tokens = x
+        return patch_tokens                  # [B, N, D]
+
+
+
+# 14. Segmentation decoder head (Segmenter-style linear decoder)
+class LinearSegmentationHead(nn.Module):
+    def __init__(self, embed_dim, num_classes, grid_size, image_size):
+        super().__init__()
+        self.grid_size = grid_size
+        self.image_size = image_size
+        self.classifier = nn.Conv2d(embed_dim, num_classes, kernel_size=1)
+
+    def forward(self, patch_tokens):
+        # patch_tokens: [B, N, D]  ->  [B, D, gh, gw]
+        B, N, D = patch_tokens.shape
+        gh = gw = self.grid_size
+        x = patch_tokens.transpose(1, 2).reshape(B, D, gh, gw)
+
+        logits = self.classifier(x)          # [B, num_classes, gh, gw]
+        logits = F.interpolate(
+            logits, size=(self.image_size, self.image_size),
+            mode="bilinear", align_corners=False,
+        )
+        return logits                        # [B, num_classes, H, W]
+
+
+class ViTSegmentation(nn.Module):
+    def __init__(self, cfg: Config):
+        super().__init__()
+        self.backbone = ViTBackbone(
+            image_size=cfg.image_size,
+            patch_size=cfg.patch_size,
+            embed_dim=cfg.embed_dim,
+            depth=cfg.depth,
+            num_heads=cfg.num_heads,
+            mlp_ratio=cfg.mlp_ratio,
+            dropout=cfg.dropout,
+            use_cls_token=True,
+        )
+        self.head = LinearSegmentationHead(
+            embed_dim=cfg.embed_dim,
+            num_classes=cfg.num_classes,
+            grid_size=self.backbone.grid_size,
+            image_size=cfg.image_size,
+        )
+
+    def forward(self, x):
+        tokens = self.backbone(x)
+        logits = self.head(tokens)
+        return logits
+
+
+
+# 15. Instantiate and sanity-check the model
+
+model = ViTSegmentation(cfg)
+
+# Use both Tesla T4 GPUs
+if torch.cuda.device_count() > 1:
+    print(f"Using {torch.cuda.device_count()} GPUs")
+    model = torch.nn.DataParallel(model)
+
+model = model.to(cfg.device)
+
+n_params = sum(
+    p.numel() for p in model.parameters()
+    if p.requires_grad
+)
+
+print(f"Trainable parameters: {n_params / 1e6:.2f}M")
+
+with torch.no_grad():
+    dummy = torch.randn(
+        2,
+        3,
+        cfg.image_size,
+        cfg.image_size,
+        device=cfg.device
+    )
+
+    out = model(dummy)
+
+print("Output shape:", out.shape)
+# [2, num_classes, image_size, image_size]
+
+
+
+
+
+# 16. Loss
+criterion = nn.CrossEntropyLoss(ignore_index=cfg.ignore_index)
+
+
+
+# 17. Segmentation metrics: pixel accuracy, mean class accuracy, mIoU
+class SegmentationMetrics:
+    def __init__(self, num_classes, ignore_index=0):
+        self.num_classes = num_classes
+        self.ignore_index = ignore_index
+        self.confusion = np.zeros((num_classes, num_classes), dtype=np.int64)
+
+    def update(self, preds, targets):
+        preds = preds.detach().cpu().numpy().reshape(-1)
+        targets = targets.detach().cpu().numpy().reshape(-1)
+        mask = targets != self.ignore_index
+        preds, targets = preds[mask], targets[mask]
+        idx = targets * self.num_classes + preds
+        cm = np.bincount(idx, minlength=self.num_classes ** 2)
+        self.confusion += cm.reshape(self.num_classes, self.num_classes)
+
+    def compute(self):
+        cm = self.confusion.astype(np.float64)
+        intersection = np.diag(cm)
+        union = cm.sum(1) + cm.sum(0) - intersection
+        iou = intersection / np.maximum(union, 1)
+        valid = union > 0
+
+        pixel_acc = intersection.sum() / max(cm.sum(), 1)
+        mean_acc = np.mean(
+            (intersection[valid] / np.maximum(cm.sum(1)[valid], 1))
+        ) if valid.any() else 0.0
+        miou = np.mean(iou[valid]) if valid.any() else 0.0
+
+        return {
+            "pixel_accuracy": float(pixel_acc),
+            "mean_class_accuracy": float(mean_acc),
+            "mIoU": float(miou),
+            "per_class_iou": iou,
+        }
+
+    def reset(self):
+        self.confusion[:] = 0
+
+
+
+# 18. Optimizer + cosine LR schedule with linear warmup
+optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
+
+steps_per_epoch = max(len(train_loader), 1)
+total_steps = steps_per_epoch * cfg.epochs
+warmup_steps = steps_per_epoch * cfg.warmup_epochs
+
+def lr_lambda(step):
+    if step < warmup_steps:
+        return step / max(1, warmup_steps)
+    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+    return 0.5 * (1 + math.cos(math.pi * progress))
+
+scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+
+
+
+
+# 19. train / eval epoch functions
+def train_one_epoch(model, loader, optimizer, scheduler, criterion, device):
+    model.train()
+    running_loss = 0.0
+    lrs = []
+    for imgs, masks in tqdm(loader, desc="train", leave=False):
+        imgs, masks = imgs.to(device), masks.to(device)
+
+        optimizer.zero_grad()
+        logits = model(imgs)
+        loss = criterion(logits, masks)
+        loss.backward()
+        optimizer.step()
+        scheduler.step()
+
+        running_loss += loss.item() * imgs.size(0)
+        lrs.append(scheduler.get_last_lr()[0])
+
+    return running_loss / len(loader.dataset), lrs
+
+
+@torch.no_grad()
+def evaluate(model, loader, criterion, device, num_classes, ignore_index):
+    model.eval()
+    running_loss = 0.0
+    metrics = SegmentationMetrics(num_classes, ignore_index)
+
+    for imgs, masks in tqdm(loader, desc="val", leave=False):
+        imgs, masks = imgs.to(device), masks.to(device)
+        logits = model(imgs)
+        loss = criterion(logits, masks)
+        running_loss += loss.item() * imgs.size(0)
+
+        preds = logits.argmax(dim=1)
+        metrics.update(preds, masks)
+
+    val_loss = running_loss / len(loader.dataset)
+    results = metrics.compute()
+    results["loss"] = val_loss
+    return results
+
+
+
+# 20. Run training
+history = {
+    "train_loss": [],
+    "val_loss": [],
+    "val_pixel_acc": [],
+    "val_mean_acc": [],
+    "val_miou": [],
+    "lr": [],
+}
+
+best_miou = -1.0
+best_state = None
+
+for epoch in range(1, cfg.epochs + 1):
+    t0 = time.time()
+
+    train_loss, lrs = train_one_epoch(model, train_loader, optimizer, scheduler, criterion, cfg.device)
+    val_results = evaluate(model, val_loader, criterion, cfg.device, cfg.num_classes, cfg.ignore_index)
+
+    history["train_loss"].append(train_loss)
+    history["val_loss"].append(val_results["loss"])
+    history["val_pixel_acc"].append(val_results["pixel_accuracy"])
+    history["val_mean_acc"].append(val_results["mean_class_accuracy"])
+    history["val_miou"].append(val_results["mIoU"])
+    history["lr"].extend(lrs)
+
+    if val_results["mIoU"] > best_miou:
+        best_miou = val_results["mIoU"]
+        best_state = copy.deepcopy(model.state_dict())
+
+    dt = time.time() - t0
+    print(
+        f"Epoch {epoch}/{cfg.epochs} | "
+        f"train_loss={train_loss:.4f} | val_loss={val_results['loss']:.4f} | "
+        f"val_mIoU={val_results['mIoU']:.4f} | pixel_acc={val_results['pixel_accuracy']:.4f} | "
+        f"{dt:.1f}s"
+    )
+
+print(f"Best val mIoU: {best_miou:.4f}")
+
+
+
+os.makedirs("checkpoints", exist_ok=True)
+
+if best_state is not None:
+    save_file(
+        best_state,
+        "checkpoints/vit_seg_best.safetensors"
+    )
+
+    print("Saved checkpoints/vit_seg_best.safetensors")
+
+
+
+
+
+# 22. Plot train/val loss, mIoU, pixel accuracy, and the LR schedule
+fig, axes = plt.subplots(2, 2, figsize=(13, 9))
+
+epochs_range = range(1, len(history["train_loss"]) + 1)
+
+axes[0, 0].plot(epochs_range, history["train_loss"], marker="o", label="train loss")
+axes[0, 0].plot(epochs_range, history["val_loss"], marker="o", label="val loss")
+axes[0, 0].set_title("Loss")
+axes[0, 0].set_xlabel("epoch")
+axes[0, 0].set_ylabel("cross-entropy loss")
+axes[0, 0].legend()
+axes[0, 0].grid(alpha=0.3)
+
+axes[0, 1].plot(epochs_range, history["val_miou"], marker="o", color="green")
+axes[0, 1].set_title("Validation mIoU")
+axes[0, 1].set_xlabel("epoch")
+axes[0, 1].set_ylabel("mIoU")
+axes[0, 1].grid(alpha=0.3)
+
+axes[1, 0].plot(epochs_range, history["val_pixel_acc"], marker="o", label="pixel accuracy")
+axes[1, 0].plot(epochs_range, history["val_mean_acc"], marker="o", label="mean class accuracy")
+axes[1, 0].set_title("Validation Accuracy")
+axes[1, 0].set_xlabel("epoch")
+axes[1, 0].legend()
+axes[1, 0].grid(alpha=0.3)
+
+axes[1, 1].plot(history["lr"], color="purple")
+axes[1, 1].set_title("Learning Rate Schedule (per step)")
+axes[1, 1].set_xlabel("step")
+axes[1, 1].set_ylabel("lr")
+axes[1, 1].grid(alpha=0.3)
+
+plt.tight_layout()
+plt.savefig("training_curves.png", dpi=150)
+plt.show()
+
+
+
+
+
+
+# 23. ADE20K color palette (151 colors incl. background) for visualization
+def ade20k_palette(num_classes=151, seed=1):
+    rng = np.random.RandomState(seed)
+    palette = rng.randint(0, 255, size=(num_classes, 3), dtype=np.uint8)
+    palette[0] = np.array([0, 0, 0], dtype=np.uint8)  # background = black
+    return palette
+
+PALETTE = ade20k_palette(cfg.num_classes)
+
+def colorize_mask(mask, palette=PALETTE):
+    mask = mask.astype(np.int64)
+    return palette[mask]
+
+
+
+ckpt_path = "checkpoints/vit_seg_best.safetensors"
+
+if os.path.exists(ckpt_path):
+    state_dict = load_file(
+        ckpt_path,
+        device=str(cfg.device)
+    )
+    model.load_state_dict(state_dict)
+    print("Loaded best checkpoint.")
+
+model.eval()
+
+
+# 25. Run inference on a handful of validation images and visualize
+@torch.no_grad()
+def run_inference(model, dataset, indices, device):
+    imgs, gts, preds = [], [], []
+    for i in indices:
+        img_t, mask_t = dataset[i]
+        logits = model(img_t.unsqueeze(0).to(device))
+        pred = logits.argmax(dim=1).squeeze(0).cpu().numpy()
+
+        # de-normalize image for display
+        img = img_t.clone()
+        for c, (m, s) in enumerate(zip(IMAGENET_MEAN, IMAGENET_STD)):
+            img[c] = img[c] * s + m
+        img = img.permute(1, 2, 0).clamp(0, 1).numpy()
+
+        imgs.append(img)
+        gts.append(mask_t.numpy())
+        preds.append(pred)
+    return imgs, gts, preds
+
+
+sample_indices = random.sample(range(len(val_dataset)), k=min(4, len(val_dataset)))
+imgs, gts, preds = run_inference(model, val_dataset, sample_indices, cfg.device)
+
+fig, axes = plt.subplots(len(sample_indices), 3, figsize=(11, 3.5 * len(sample_indices)))
+if len(sample_indices) == 1:
+    axes = axes[None, :]
+
+for row, (img, gt, pred) in enumerate(zip(imgs, gts, preds)):
+    axes[row, 0].imshow(img)
+    axes[row, 0].set_title("Image")
+    axes[row, 0].axis("off")
+
+    axes[row, 1].imshow(colorize_mask(gt))
+    axes[row, 1].set_title("Ground Truth")
+    axes[row, 1].axis("off")
+
+    axes[row, 2].imshow(colorize_mask(pred))
+    axes[row, 2].set_title("Prediction")
+    axes[row, 2].axis("off")
+
+plt.tight_layout()
+plt.savefig("inference_samples.png", dpi=150)
+plt.show()
+
+
+
+# Config — export as JSON
+import json
+
+config_dict = {
+    "data": {
+        "dataset": "scene_parse_150 (ADE20K, 150 classes)",
+        "image_size": 224,
+        "patch_size": 16,
+        "num_classes": 151,
+        "ignore_index": 0
+    },
+    "model": {
+        "architecture": "ViT-Segmentation (first-principles backbone + linear decoder head)",
+        "embed_dim": 384,
+        "depth": 8,
+        "num_heads": 6,
+        "mlp_ratio": 4.0,
+        "dropout": 0.1
+    },
+    "training": {
+        "batch_size": 8,
+        "epochs": 5,
+        "lr": 3e-4,
+        "weight_decay": 0.05,
+        "warmup_epochs": 1,
+        "num_workers": 2,
+        "subset_size": None,
+        "device": "cuda if available else cpu"
+    },
+    "evolution_search": {
+        "population_size": 6,
+        "generations": 4,
+        "elite_k": 2,
+        "steps_per_candidate": 30,
+        "subset_size": 64,
+        "batch_size": 4
+    }
+}
+
+with open("config.json", "w") as f:
+    json.dump(config_dict, f, indent=2)
+
+
+# README.md — generated for this notebook, includes result images
+readme_content = """
+# ViT Semantic Segmentation on ADE20K
+
+A Vision Transformer built from first principles (patch embedding, multi-head
+self-attention, MLP, transformer encoder blocks) adapted for semantic
+segmentation on the ADE20K dataset (150 classes), with training, evaluation,
+benchmarking against published ViT segmentation models, and an evolutionary
+hyperparameter search.
+
+## Contents
+
+- `vit-ade20k-segmentation.ipynb` — the full notebook
+- `config.json` — model/training/search configuration
+- `checkpoints/vit_seg_best.pt` — best model checkpoint (created after training)
+
+## Configuration
+
+See [`config.json`](config.json) for the exact data, model, training, and
+evolutionary-search settings used to produce the results below.
+
+## Results
+
+### Training curves
+
+get_ipython().getoutput("[Training curves](training_curves.png)")
+
+Train/val loss, validation mIoU, validation pixel & mean-class accuracy, and
+the learning-rate schedule.
+
+### Inference samples
+
+get_ipython().getoutput("[Inference samples](inference_samples.png)")
+
+Image / ground truth / prediction, side by side, for a few validation images.
+
+### Benchmark comparison
+
+get_ipython().getoutput("[Benchmark comparison](benchmark_comparison.png)")
+
+mIoU compared against published ViT-based segmentation models (SETR,
+Segmenter, SegFormer, DPT) on ADE20K val, plus an accuracy-vs-parameters
+scatter.
+
+### Evolutionary search
+
+get_ipython().getoutput("[Evolution fitness](evolution_fitness.png)")
+
+Best and mean fitness (proxy mIoU) per generation from the evolutionary
+hyperparameter search.
+
+## How to reproduce
+
+1. Install dependencies (see the first cell of the notebook).
+2. Run all cells top to bottom.
+3. Images (`training_curves.png`, `inference_samples.png`,
+   `benchmark_comparison.png`, `evolution_fitness.png`) and `config.json`
+   are written to the working directory as you go.
+
+## Notes
+
+- Published benchmark numbers are commonly-cited approximations from the
+  original papers — verify exact figures before citing them elsewhere.
+- Demo defaults (small `embed_dim`/`depth`, few epochs) are set to run
+  quickly; scale up for real training.
+"""
+
+with open("README.md", "w") as f:
+    f.write(readme_content)
+
+
+get_ipython().run_line_magic("ls", "")
+
+
+import shutil
+
+os.makedirs("/kaggle/working/vit", exist_ok=True)
+
+for item in os.listdir("/kaggle/working"):
+    if item != "vit":
+        shutil.move(
+            os.path.join("/kaggle/working", item),
+            "/kaggle/working/vit"
+        )
+
+
+from huggingface_hub import HfApi
+
+api = HfApi(token="hf_xxxx")
+
+api.create_repo(
+    repo_id="aijadugar/ViT-ADE20K-Segmentation",
+    exist_ok=True,
+)
+
+api.upload_folder(
+    repo_id="aijadugar/ViT-ADE20K-Segmentation",
+    folder_path="/kaggle/working/vit",
+    commit_message="Implemented Vision Transformer architechture from scratch!",
+    repo_type="model",
+)
+
+
+

README.md

@@ -0,0 +1,64 @@
+
+# ViT Semantic Segmentation on ADE20K
+
+A Vision Transformer built from first principles (patch embedding, multi-head
+self-attention, MLP, transformer encoder blocks) adapted for semantic
+segmentation on the ADE20K dataset (150 classes), with training, evaluation,
+benchmarking against published ViT segmentation models, and an evolutionary
+hyperparameter search.
+
+## Contents
+
+- `vit-ade20k-segmentation.ipynb` — the full notebook
+- `config.json` — model/training/search configuration
+- `checkpoints/vit_seg_best.pt` — best model checkpoint (created after training)
+
+## Configuration
+
+See [`config.json`](config.json) for the exact data, model, training, and
+evolutionary-search settings used to produce the results below.
+
+## Results
+
+### Training curves
+
+![Training curves](training_curves.png)
+
+Train/val loss, validation mIoU, validation pixel & mean-class accuracy, and
+the learning-rate schedule.
+
+### Inference samples
+
+![Inference samples](inference_samples.png)
+
+Image / ground truth / prediction, side by side, for a few validation images.
+
+### Benchmark comparison
+
+![Benchmark comparison](benchmark_comparison.png)
+
+mIoU compared against published ViT-based segmentation models (SETR,
+Segmenter, SegFormer, DPT) on ADE20K val, plus an accuracy-vs-parameters
+scatter.
+
+### Evolutionary search
+
+![Evolution fitness](evolution_fitness.png)
+
+Best and mean fitness (proxy mIoU) per generation from the evolutionary
+hyperparameter search.
+
+## How to reproduce
+
+1. Install dependencies (see the first cell of the notebook).
+2. Run all cells top to bottom.
+3. Images (`training_curves.png`, `inference_samples.png`,
+   `benchmark_comparison.png`, `evolution_fitness.png`) and `config.json`
+   are written to the working directory as you go.
+
+## Notes
+
+- Published benchmark numbers are commonly-cited approximations from the
+  original papers — verify exact figures before citing them elsewhere.
+- Demo defaults (small `embed_dim`/`depth`, few epochs) are set to run
+  quickly; scale up for real training.

checkpoints/vit_seg_best.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03051ca030569c0f53f561b3ab874eb7326bde8290ee44ceb6d4ce55faeecbde
+size 58515052

config.json

@@ -0,0 +1,35 @@
+{
+  "data": {
+    "dataset": "scene_parse_150 (ADE20K, 150 classes)",
+    "image_size": 224,
+    "patch_size": 16,
+    "num_classes": 151,
+    "ignore_index": 0
+  },
+  "model": {
+    "architecture": "ViT-Segmentation (first-principles backbone + linear decoder head)",
+    "embed_dim": 384,
+    "depth": 8,
+    "num_heads": 6,
+    "mlp_ratio": 4.0,
+    "dropout": 0.1
+  },
+  "training": {
+    "batch_size": 8,
+    "epochs": 5,
+    "lr": 0.0003,
+    "weight_decay": 0.05,
+    "warmup_epochs": 1,
+    "num_workers": 2,
+    "subset_size": null,
+    "device": "cuda if available else cpu"
+  },
+  "evolution_search": {
+    "population_size": 6,
+    "generations": 4,
+    "elite_k": 2,
+    "steps_per_candidate": 30,
+    "subset_size": 64,
+    "batch_size": 4
+  }
+}