{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.13.15","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"accelerator":"GPU"},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"2acf793f-e274-491d-a802-04fe1a8412a4","cell_type":"markdown","source":"# CIFAR-10 Challenge: RouNet trained from scratch\n\nKaggle competition: https://www.kaggle.com/c/cifar-10\n\n**Requirements covered**\n- **Random initialization only.** No pre-trained weights are used; convolutions use Kaiming-normal init.\n- **Advanced optimizer:** AdamW with decoupled weight decay (BN parameters and biases are excluded from decay).\n- **LR schedule:** linear warmup (5 epochs), then cosine annealing to 1% of the peak LR, updated every step.\n- **Advanced augmentation:** CutMix (p = 0.5, switched off for the final 10 epochs) plus Cutout on the remaining batches, on top of random crop and horizontal flip. All augmentation runs on the GPU.\n- **Custom network (RouNet):** three downsampling stages, each with a residual block that has Squeeze-and-Excitation channel attention, and a dual (average + max) pooling head.\n- **Extras:** label smoothing 0.1, weight EMA, mixed precision, and horizontal-flip test-time augmentation.\n\n**Kaggle settings:** Accelerator = **GPU T4 x2** (the Kaggle PyTorch build lacks P100 kernels). Internet = On (only needed if `py7zr` must be installed).\n\n**Workflow:** run once with `val_size=5000` to monitor validation accuracy, then set `val_size=0` for the final submission run.\n","metadata":{}},{"id":"74120d35-8f1d-4050-9ef7-846def56591a","cell_type":"markdown","source":"## Setup & config","metadata":{}},{"id":"f11ca905-55dc-4b53-9238-e491bef368a3","cell_type":"code","source":"import os, math, time, shutil, subprocess, sys, glob, random\nimport numpy as np, pandas as pd, torch, torch.nn as nn, torch.nn.functional as F\nfrom PIL import Image\nfrom concurrent.futures import ThreadPoolExecutor\n\nDATA_DIR = os.environ.get(\"DATA_DIR\", \"/kaggle/input/cifar-10\")\nWORK_DIR = os.environ.get(\"WORK_DIR\", \"/tmp/cifar\")\nOUT_DIR  = os.environ.get(\"OUT_DIR\", \"/kaggle/working\")\n\nCFG = dict(\n    epochs=100,          # total epochs\n    warmup_epochs=5,     # linear warmup, then cosine annealing\n    batch_size=512,\n    lr=2e-3,             # peak AdamW learning rate\n    min_lr_ratio=0.01,   # cosine floor = lr * 0.01\n    weight_decay=5e-2,   # decoupled weight decay (AdamW); not applied to BN/bias\n    label_smoothing=0.1,\n    cutmix_p=0.5,        # probability of applying CutMix to a batch\n    cutmix_off_epochs=10,# disable CutMix for the final N epochs\n    cutout_size=8,       # Cutout patch on non-CutMix batches (0 = off)\n    ema_decay=0.999,\n    val_size=0,       # set to 0 for the final run (train on all 50k)\n    seed=0,\n)\n\ndef seed_all(s):\n    random.seed(s); np.random.seed(s); torch.manual_seed(s); torch.cuda.manual_seed_all(s)\nseed_all(CFG[\"seed\"])\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nn_gpus = torch.cuda.device_count()\ntorch.backends.cudnn.benchmark = True\nprint(f\"device={device}, GPUs={n_gpus}, torch={torch.__version__}\")\nprint(CFG)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-09T11:18:27.824717Z","iopub.execute_input":"2026-10-09T11:18:27.825013Z","execution_failed":"2026-10-09T11:18:33.478Z"}},"outputs":[],"execution_count":null},{"id":"7d21631f-bddc-4964-9c65-39f6c4901f16","cell_type":"markdown","source":"## Extract & load data","metadata":{}},{"id":"8a5acaf4-1eba-4eb7-a257-537d427b3e44","cell_type":"code","source":"# Locate the competition data wherever Kaggle mounted it\nhits = glob.glob(\"/kaggle/input/**/trainLabels.csv\", recursive=True)\nassert hits, \"cifar-10 data not attached - add it via Input > Add Input\"\nDATA_DIR = os.path.dirname(hits[0])\nprint(\"DATA_DIR =\", DATA_DIR, os.listdir(DATA_DIR))\n\ndef extract(archive):\n    os.makedirs(WORK_DIR, exist_ok=True)\n    if shutil.which(\"7z\"):\n        r = subprocess.run([\"7z\", \"x\", \"-y\", archive, f\"-o{WORK_DIR}\"], capture_output=True, text=True)\n        if r.returncode == 0:\n            return\n        print(\"7z failed:\\n\", r.stdout[-2000:], r.stderr[-2000:], \"\\nfalling back to py7zr\")\n    try:\n        import py7zr\n    except ImportError:\n        subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"py7zr\"], check=True)\n        import py7zr\n    with py7zr.SevenZipFile(archive) as z:\n        z.extractall(WORK_DIR)\n\ndef load_folder(folder):\n    n = len(glob.glob(f\"{folder}/*.png\"))\n    read = lambda i: np.asarray(Image.open(f\"{folder}/{i}.png\").convert(\"RGB\"))\n    with ThreadPoolExecutor(16) as ex:\n        imgs = list(ex.map(read, range(1, n + 1)))   # files are named 1.png ... n.png\n    return np.stack(imgs)                            # (n, 32, 32, 3) uint8\n\nt0 = time.time()\nfor name in [\"train\", \"test\"]:\n    if not os.path.isdir(f\"{WORK_DIR}/{name}\"):\n        extract(f\"{DATA_DIR}/{name}.7z\")\n\nfor name in [\"train\", \"test\"]:\n    if not os.path.isdir(f\"{WORK_DIR}/{name}\"):\n        extract(f\"{DATA_DIR}/{name}.7z\")\n        \nX_all = load_folder(f\"{WORK_DIR}/train\")\nX_test = load_folder(f\"{WORK_DIR}/test\")\nlabels = pd.read_csv(f\"{DATA_DIR}/trainLabels.csv\").sort_values(\"id\")\nCLASSES = sorted(labels[\"label\"].unique())\ny_all = labels[\"label\"].map({c: i for i, c in enumerate(CLASSES)}).values\nprint(f\"train {X_all.shape}, test {X_test.shape}, classes {CLASSES}, load {time.time()-t0:.0f}s\")\n\nperm = np.random.permutation(len(X_all))\nval_idx, tr_idx = perm[:CFG[\"val_size\"]], perm[CFG[\"val_size\"]:]\n\n# Keep the training set on the GPU as uint8; all augmentation runs on the GPU.\nX_tr = torch.from_numpy(X_all[tr_idx]).permute(0, 3, 1, 2).contiguous().to(device)\ny_tr = torch.from_numpy(y_all[tr_idx]).long().to(device)\nX_val = torch.from_numpy(X_all[val_idx]).permute(0, 3, 1, 2).contiguous()\ny_val = torch.from_numpy(y_all[val_idx]).long()\n\nMEAN = torch.tensor([0.4914, 0.4822, 0.4465], device=device).view(1, 3, 1, 1) * 255\nSTD  = torch.tensor([0.2470, 0.2435, 0.2616], device=device).view(1, 3, 1, 1) * 255\ndef normalize(x_uint8):\n    return ((x_uint8.to(device, non_blocking=True).float() - MEAN) / STD).contiguous(memory_format=torch.channels_last)","metadata":{"trusted":true,"execution":{"execution_failed":"2026-10-09T11:18:33.479Z"}},"outputs":[],"execution_count":null},{"id":"257413fc-8746-4756-9bd4-0bf7fc33e300","cell_type":"markdown","source":"## GPU augmentation","metadata":{}},{"id":"baaa61bf-47b6-4f6a-a243-4ae8727a7c50","cell_type":"code","source":"def random_crop_flip(x, pad=4):\n    N, C, H, W = x.shape\n    x = F.pad(x, (pad, pad, pad, pad), mode=\"reflect\")\n    i = torch.randint(0, 2 * pad + 1, (N,), device=x.device)\n    j = torch.randint(0, 2 * pad + 1, (N,), device=x.device)\n    ar = torch.arange(H, device=x.device)\n    rows, cols = (i[:, None] + ar), (j[:, None] + ar)\n    x = x[torch.arange(N, device=x.device)[:, None, None], :, rows[:, :, None], cols[:, None, :]]  # N,H,W,C\n    x = x.permute(0, 3, 1, 2)\n    flip = torch.rand(N, device=x.device) < 0.5\n    return torch.where(flip[:, None, None, None], x.flip(3), x)\n\ndef rand_box(H, W, lam):\n    cut = math.sqrt(1 - lam)\n    ch, cw = int(H * cut), int(W * cut)\n    cy, cx = np.random.randint(H), np.random.randint(W)\n    y1, y2 = np.clip(cy - ch // 2, 0, H), np.clip(cy + ch // 2, 0, H)\n    x1, x2 = np.clip(cx - cw // 2, 0, W), np.clip(cx + cw // 2, 0, W)\n    return y1, y2, x1, x2\n\ndef cutmix(x, y):\n    lam = np.random.beta(1.0, 1.0)\n    idx = torch.randperm(x.size(0), device=x.device)\n    y1, y2, x1, x2 = rand_box(x.size(2), x.size(3), lam)\n    x = x.clone()\n    x[:, :, y1:y2, x1:x2] = x[idx, :, y1:y2, x1:x2]\n    lam = 1 - (y2 - y1) * (x2 - x1) / (x.size(2) * x.size(3))   # exact area ratio\n    return x, y, y[idx], lam\n\ndef cutout(x, size):\n    N, _, H, W = x.shape\n    cy = torch.randint(0, H, (N, 1, 1), device=x.device)\n    cx = torch.randint(0, W, (N, 1, 1), device=x.device)\n    yy = torch.arange(H, device=x.device)[None, :, None]\n    xx = torch.arange(W, device=x.device)[None, None, :]\n    mask = ((yy - cy).abs() <= size // 2) & ((xx - cx).abs() <= size // 2)\n    return x * (~mask)[:, None].to(x.dtype)","metadata":{},"outputs":[],"execution_count":null},{"id":"7b5bf6e1-2e85-4b8a-8e90-c969d40926e7","cell_type":"markdown","source":"## Custom network","metadata":{}},{"id":"a1f83361-ae50-46c5-a08a-5582da2ea795","cell_type":"code","source":"class ConvBN(nn.Sequential):\n    def __init__(self, cin, cout):\n        super().__init__(nn.Conv2d(cin, cout, 3, padding=1, bias=False),\n                         nn.BatchNorm2d(cout), nn.ReLU(inplace=True))\n\nclass SE(nn.Module):\n    \"\"\"Squeeze-and-Excitation: re-weights channels using global context.\"\"\"\n    def __init__(self, c, r=8):\n        super().__init__()\n        self.fc = nn.Sequential(nn.Linear(c, c // r), nn.ReLU(inplace=True), nn.Linear(c // r, c), nn.Sigmoid())\n    def forward(self, x):\n        return x * self.fc(x.mean((2, 3)))[:, :, None, None]\n\nclass ResSE(nn.Module):\n    def __init__(self, c):\n        super().__init__()\n        self.body = nn.Sequential(ConvBN(c, c), ConvBN(c, c), SE(c))\n    def forward(self, x):\n        return x + self.body(x)\n\nclass Stage(nn.Sequential):\n    \"\"\"Widen -> downsample -> residual SE block.\"\"\"\n    def __init__(self, cin, cout):\n        super().__init__(ConvBN(cin, cout), nn.MaxPool2d(2), ResSE(cout))\n\nclass DualPoolHead(nn.Module):\n    \"\"\"Concatenates global average and global max pooling before the classifier.\"\"\"\n    def __init__(self, c, n_classes, p=0.2):\n        super().__init__()\n        self.drop, self.fc = nn.Dropout(p), nn.Linear(2 * c, n_classes)\n    def forward(self, x):\n        x = torch.cat([F.adaptive_avg_pool2d(x, 1), F.adaptive_max_pool2d(x, 1)], 1).flatten(1)\n        return self.fc(self.drop(x))\n\nclass RouNet(nn.Module):\n    def __init__(self, widths=(64, 128, 256, 512), n_classes=10):\n        super().__init__()\n        w0, w1, w2, w3 = widths\n        self.net = nn.Sequential(\n            ConvBN(3, w0),        # 32x32\n            Stage(w0, w1),        # 16x16\n            Stage(w1, w2),        # 8x8\n            Stage(w2, w3),        # 4x4\n            DualPoolHead(w3, n_classes))\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                nn.init.kaiming_normal_(m.weight, mode=\"fan_out\", nonlinearity=\"relu\")  # random init only\n    def forward(self, x):\n        return self.net(x)\n\nmodel = RouNet().to(device).to(memory_format=torch.channels_last)\nprint(f\"RouNet parameters: {sum(p.numel() for p in model.parameters())/1e6:.2f}M\")","metadata":{},"outputs":[],"execution_count":null},{"id":"808f0e1c-9502-467c-af99-fcb43db424ca","cell_type":"markdown","source":"## Optimizer, schedule, EMA","metadata":{}},{"id":"1a41faeb-51af-45d3-bcc5-da499cf73cdc","cell_type":"code","source":"decay, no_decay = [], []\nfor n_, p in model.named_parameters():\n    (no_decay if p.ndim <= 1 else decay).append(p)       # no weight decay on BN params and biases\nopt = torch.optim.AdamW([{\"params\": decay, \"weight_decay\": CFG[\"weight_decay\"]},\n                         {\"params\": no_decay, \"weight_decay\": 0.0}],\n                        lr=CFG[\"lr\"], betas=(0.9, 0.999))\n\nsteps_per_epoch = math.ceil(len(X_tr) / CFG[\"batch_size\"])\ntotal_steps = CFG[\"epochs\"] * steps_per_epoch\nwarmup_steps = CFG[\"warmup_epochs\"] * steps_per_epoch\ndef lr_lambda(step):\n    if step < warmup_steps:\n        return (step + 1) / warmup_steps\n    t = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n    return CFG[\"min_lr_ratio\"] + (1 - CFG[\"min_lr_ratio\"]) * 0.5 * (1 + math.cos(math.pi * t))\nsched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)\n\nema = torch.optim.swa_utils.AveragedModel(\n    model, multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(CFG[\"ema_decay\"]), use_buffers=True)\n\ntrain_model = nn.DataParallel(model) if n_gpus > 1 else model   # uses both T4s if available\nuse_amp = device == \"cuda\"\nscaler = torch.amp.GradScaler(\"cuda\", enabled=use_amp)\ncriterion = nn.CrossEntropyLoss(label_smoothing=CFG[\"label_smoothing\"])\n\n@torch.no_grad()\ndef predict(net, X_uint8, bs=1024, tta=True):\n    net.eval(); out = []\n    for k in range(0, len(X_uint8), bs):\n        x = normalize(X_uint8[k:k + bs])\n        with torch.autocast(device_type=\"cuda\", dtype=torch.float16, enabled=use_amp):\n            p = net(x).float().softmax(1)\n            if tta:\n                p = (p + net(x.flip(3)).float().softmax(1)) / 2   # horizontal-flip TTA\n        out.append(p.cpu())\n    return torch.cat(out)","metadata":{},"outputs":[],"execution_count":null},{"id":"71fb811a-bd3e-4466-afc3-2c33653cc4bc","cell_type":"markdown","source":"## Training loop","metadata":{}},{"id":"0cd8cd43-be04-4ce5-bf4d-9c9a5303a272","cell_type":"code","source":"history = []\nstep = 0\nfor epoch in range(CFG[\"epochs\"]):\n    train_model.train()\n    t0, loss_sum, n_seen = time.time(), 0.0, 0\n    use_cutmix = epoch < CFG[\"epochs\"] - CFG[\"cutmix_off_epochs\"]\n    order = torch.randperm(len(X_tr), device=device)\n    for k in range(0, len(X_tr), CFG[\"batch_size\"]):\n        idx = order[k:k + CFG[\"batch_size\"]]\n        x, y = random_crop_flip(normalize(X_tr[idx])), y_tr[idx]\n        with torch.autocast(device_type=\"cuda\", dtype=torch.float16, enabled=use_amp):\n            if use_cutmix and np.random.rand() < CFG[\"cutmix_p\"]:\n                x, ya, yb, lam = cutmix(x, y)\n                logits = train_model(x.contiguous(memory_format=torch.channels_last))\n                loss = lam * criterion(logits, ya) + (1 - lam) * criterion(logits, yb)\n            else:\n                if CFG[\"cutout_size\"] > 0:\n                    x = cutout(x, CFG[\"cutout_size\"])\n                logits = train_model(x.contiguous(memory_format=torch.channels_last))\n                loss = criterion(logits, y)\n        opt.zero_grad(set_to_none=True)\n        scaler.scale(loss).backward()\n        scaler.step(opt); scaler.update(); sched.step(); step += 1\n        ema.update_parameters(model)\n        loss_sum += loss.item() * len(idx); n_seen += len(idx)\n\n    rec = dict(epoch=epoch + 1, loss=loss_sum / n_seen, lr=sched.get_last_lr()[0], time=time.time() - t0)\n    if len(X_val) and ((epoch + 1) % 5 == 0 or epoch + 1 == CFG[\"epochs\"]):\n        rec[\"val_acc\"] = (predict(model, X_val, tta=False).argmax(1) == y_val).float().mean().item()\n        rec[\"val_acc_ema_tta\"] = (predict(ema.module, X_val).argmax(1) == y_val).float().mean().item()\n    history.append(rec)\n    print(\" | \".join(f\"{k}={v:.4f}\" if isinstance(v, float) else f\"{k}={v}\" for k, v in rec.items()))\n\npd.DataFrame(history).to_csv(f\"{OUT_DIR}/history.csv\", index=False)\ntorch.save(ema.module.state_dict(), f\"{OUT_DIR}/rounet_ema.pt\")","metadata":{},"outputs":[],"execution_count":null},{"id":"292ee9cf-605d-4dd5-be91-ff929d24846c","cell_type":"markdown","source":"## Predict test set & write submission","metadata":{}},{"id":"04b58768-53af-40d8-9422-66d3a8f04108","cell_type":"code","source":"probs = predict(ema.module, torch.from_numpy(X_test).permute(0, 3, 1, 2).contiguous())\nsub = pd.DataFrame({\"id\": np.arange(1, len(X_test) + 1), \"label\": np.array(CLASSES)[probs.argmax(1).numpy()]})\nsub.to_csv(f\"{OUT_DIR}/submission.csv\", index=False)\nprint(sub.head(), sub[\"label\"].value_counts(), sep=\"\\n\")","metadata":{},"outputs":[],"execution_count":null},{"id":"3c4a5561-b9ab-4cc3-9ae4-87c356cbb8c7","cell_type":"markdown","source":"## Plots for the report","metadata":{}},{"id":"e17b8cf9-530a-4722-bbf4-14dc2ea2c65a","cell_type":"code","source":"import matplotlib.pyplot as plt\nh = pd.DataFrame(history)\nlrs = [CFG[\"lr\"] * lr_lambda(s) for s in range(total_steps)]\nfig, ax = plt.subplots(1, 3, figsize=(15, 4))\nax[0].plot(np.arange(total_steps) / steps_per_epoch, lrs); ax[0].set(title=\"LR: warmup + cosine\", xlabel=\"epoch\")\nax[1].plot(h.epoch, h.loss); ax[1].set(title=\"Training loss\", xlabel=\"epoch\")\nif \"val_acc\" in h:\n    v = h.dropna(subset=[\"val_acc\"])\n    ax[2].plot(v.epoch, v.val_acc, \"o-\", label=\"model\"); ax[2].plot(v.epoch, v.val_acc_ema_tta, \"o-\", label=\"EMA + TTA\")\n    ax[2].legend(); ax[2].set(title=\"Validation accuracy\", xlabel=\"epoch\")\nplt.tight_layout(); plt.savefig(f\"{OUT_DIR}/training_curves.png\", dpi=150)","metadata":{},"outputs":[],"execution_count":null}]}