# 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", )