{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# %% [markdown]\n# # CIFAR-10 from scratch: ResNet-style ConvNet (Kaggle GPU)\n# * Randomly initialised ConvNet, **no pretrained weights, no transfer learning**\n# * Official CIFAR-10 training images split **80% train / 20% validation** (stratified, fixed seed)\n# * Validation set is used only for monitoring / model selection, never for gradient updates\n# * Optional final stage: retrain the *selected* configuration on all 50,000 training images\n# * Kaggle test images are only used for inference\n#\n# **Settings needed:** Accelerator = GPU. Internet = ON only if the `.7z` archives must be extracted and `py7zr` is not installed.\n\n# %% [markdown]\n# ## 1. Configuration\n\n# %%\nimport os, sys, time, glob, shutil, random, subprocess, math\nfrom concurrent.futures import ThreadPoolExecutor\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import TensorDataset, DataLoader\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom sklearn.model_selection import train_test_split\nfrom PIL import Image\n\n# ------------------------- main hyper-parameters (edit here) -------------------------\nSEED          = 42\nBATCH_SIZE    = 128\nEPOCHS        = 40          # epochs of the final 80/20 run (and of the optional full retrain)\nLEARNING_RATE = 0.1         # peak LR for SGD (batch 128)\nWEIGHT_DECAY  = 5e-4\nMOMENTUM      = 0.9\nWARMUP_EPOCHS = 2\nNUM_WORKERS   = 2\nVAL_FRACTION  = 0.2         # -> 40,000 train / 10,000 validation\n\nRUN_EXPERIMENTS = True      # small augmentation study on the validation set\nEXP_EPOCHS      = 15        # epochs per experiment (short runs to save GPU time)\nDEFAULT_AUG     = \"C_cutmix_ls\"   # used only if RUN_EXPERIMENTS = False\nRETRAIN_FULL    = True      # retrain chosen config on all 50k images after model selection\n# --------------------------------------------------------------------------------------\n\nIS_KAGGLE = os.path.exists(\"/kaggle/working\")\nWORK_DIR  = \"/kaggle/working\" if IS_KAGGLE else \"./working\"\nDATA_ROOT = \"/kaggle/input\"\nEXTRACT_DIR = \"/kaggle/temp/cifar10\" if os.path.exists(\"/kaggle/temp\") else \"/tmp/cifar10\"\nSUBMISSION_PATH = os.path.join(WORK_DIR, \"submission.csv\")\nos.makedirs(WORK_DIR, exist_ok=True)\n\ndevice  = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nUSE_AMP = device.type == \"cuda\"          # mixed precision for speed on GPU\nprint(\"Device:\", device, \"|\", torch.cuda.get_device_name(0) if USE_AMP else \"CPU (this will be very slow)\")\n\n\ndef set_seed(seed=SEED):\n    \"\"\"Seed Python, NumPy, PyTorch. cudnn.benchmark is kept ON for speed, so GPU runs are\n    reproducible in data split/initialisation but not bit-for-bit identical.\"\"\"\n    random.seed(seed); np.random.seed(seed)\n    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.benchmark = True\n\nset_seed()\n\n# %% [markdown]\n# ## 2. Locate / extract the Kaggle data\n# The competition provides `train.7z`, `test.7z`, `trainLabels.csv`, `sampleSubmission.csv`.\n# If image folders already exist in `/kaggle/input` they are used; otherwise the archives are extracted to a temp folder (not to the output folder).\n\n# %%\ndef locate(*names):\n    \"\"\"Case-insensitive search for any of the given file names under /kaggle/input.\"\"\"\n    wanted = {n.lower() for n in names}\n    for root, _, files in os.walk(DATA_ROOT):\n        for f in files:\n            if f.lower() in wanted:\n                return os.path.join(root, f)\n    return None\n\nLABELS_CSV = locate(\"trainLabels.csv\", \"train_labels.csv\")\nSAMPLE_CSV = locate(\"sampleSubmission.csv\", \"sample_submission.csv\")\nif not (LABELS_CSV and SAMPLE_CSV):\n    for root, dirs, files in os.walk(DATA_ROOT):          # show what IS attached, to help debugging\n        print(root, \"->\", dirs[:5], files[:8])\n        if root.count(\"/\") > 5: break\nassert LABELS_CSV and SAMPLE_CSV, \"Could not find trainLabels.csv / sampleSubmission.csv - attach the CIFAR-10 competition data (Add Input).\"\nCOMP_DIR = os.path.dirname(LABELS_CSV)\nprint(\"Competition folder:\", COMP_DIR, \"->\", sorted(os.listdir(COMP_DIR)))\n\n\ndef extract_7z(archive, out_dir):\n    \"\"\"Prefer the native multi-threaded 7z binary (much faster than py7zr for 300k tiny PNGs).\"\"\"\n    os.makedirs(out_dir, exist_ok=True)\n    exe = shutil.which(\"7z\") or shutil.which(\"7zz\") or shutil.which(\"7za\") or shutil.which(\"7zr\")\n    if exe is None:                                    # try to install it (needs Internet ON)\n        try:\n            subprocess.run([\"apt-get\", \"update\", \"-qq\"], check=True)\n            subprocess.run([\"apt-get\", \"install\", \"-y\", \"-qq\", \"p7zip-full\"], check=True)\n            exe = shutil.which(\"7z\")\n        except Exception as e:\n            print(\"apt install of p7zip-full failed:\", e)\n    if exe is not None:\n        print(\"Extracting with\", exe)\n        proc = subprocess.Popen([exe, \"x\", archive, f\"-o{out_dir}\", \"-y\"],\n                                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)\n        t0 = time.time()\n        while proc.poll() is None:                      # progress report every 20 s\n            time.sleep(20)\n            n = sum(len(f) for _, _, f in os.walk(out_dir))\n            print(f\"  {n} files extracted, {time.time()-t0:.0f}s elapsed\", flush=True)\n        assert proc.returncode == 0, f\"7z failed with code {proc.returncode}\"\n        return\n    print(\"Falling back to py7zr (slow for 300k files)\")\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, mode=\"r\") as z:\n        z.extractall(out_dir)\n\n\ndef find_image_dir(root, probe=\"1.png\"):\n    for dirpath, _, files in os.walk(root):\n        if probe in files:\n            return dirpath\n    return None\n\n\ndef count_pngs(d):\n    return sum(f.endswith(\".png\") for f in os.listdir(d))\n\n\ndef prepare_image_dir(split, expected):\n    \"\"\"Return a folder with >= `expected` <id>.png files for split in {'train','test'}.\n    Uses images already present in /kaggle/input or a previous extraction; otherwise extracts the .7z.\n    A partially extracted folder (count too low) is deleted and extracted again.\"\"\"\n    for root in (os.path.join(COMP_DIR, split), os.path.join(EXTRACT_DIR, split)):\n        d = find_image_dir(root)\n        if d is not None and count_pngs(d) >= expected:\n            return d\n    archive = os.path.join(COMP_DIR, f\"{split}.7z\")\n    assert os.path.exists(archive), f\"{archive} not found and no extracted {split} images available\"\n    out = os.path.join(EXTRACT_DIR, split)\n    shutil.rmtree(out, ignore_errors=True)\n    print(f\"Extracting {archive} -> {out} ...\")\n    t0 = time.time(); extract_7z(archive, out)\n    print(f\"  done in {time.time()-t0:.0f}s\")\n    d = find_image_dir(out)\n    assert d is not None and count_pngs(d) >= expected, f\"extraction incomplete for {split}\"\n    return d\n\n\ndef load_images(img_dir, ids):\n    \"\"\"Load PNGs (32x32 RGB) into one uint8 tensor of shape (N, 32, 32, 3), printing progress for big sets.\"\"\"\n    def _load(i):\n        with Image.open(os.path.join(img_dir, f\"{i}.png\")) as im:\n            return np.asarray(im.convert(\"RGB\"))\n    out, chunk = [], 20000\n    with ThreadPoolExecutor(max_workers=8) as ex:\n        for s in range(0, len(ids), chunk):\n            out.extend(ex.map(_load, ids[s:s + chunk]))\n            if len(ids) > chunk:\n                print(f\"  loaded {min(s + chunk, len(ids))}/{len(ids)} images\", flush=True)\n    return torch.from_numpy(np.stack(out))\n\n# %% [markdown]\n# ## 3. Datasets: 80/20 split, normalisation statistics from the training part only\n\n# %%\ndef get_datasets():\n    labels = pd.read_csv(LABELS_CSV)                              # columns: id, label (class names)\n    classes = sorted(labels[\"label\"].unique())\n    assert len(classes) == 10, classes\n    class_to_idx = {c: i for i, c in enumerate(classes)}\n    y_all = torch.tensor(labels[\"label\"].map(class_to_idx).values, dtype=torch.long)\n\n    train_dir = prepare_image_dir(\"train\", 50000)\n    x_all = load_images(train_dir, labels[\"id\"].values)\n    assert len(x_all) == 50000, f\"expected 50,000 training images, got {len(x_all)}\"\n\n    # Stratified, seeded 80/20 split -> exactly 40,000 / 10,000\n    idx_tr, idx_val = train_test_split(np.arange(len(y_all)), test_size=VAL_FRACTION,\n                                       random_state=SEED, stratify=y_all.numpy())\n    assert len(set(idx_tr) & set(idx_val)) == 0, \"train/val overlap!\"\n    idx_tr, idx_val = torch.tensor(idx_tr), torch.tensor(idx_val)\n    data = dict(classes=classes, x_all=x_all, y_all=y_all,\n                x_train=x_all[idx_tr], y_train=y_all[idx_tr],\n                x_val=x_all[idx_val],  y_val=y_all[idx_val],\n                idx_train=idx_tr, idx_val=idx_val)\n\n    # Normalisation statistics computed on the TRAIN split only (no validation information)\n    xf = data[\"x_train\"].float().div(255)\n    data[\"mean\"] = xf.mean(dim=(0, 1, 2)).view(1, 3, 1, 1).to(device)\n    data[\"std\"]  = xf.std(dim=(0, 1, 2)).view(1, 3, 1, 1).to(device)\n    return data\n\n\ndef get_dataloaders(x_train, y_train, x_val, y_val):\n    \"\"\"Loaders yield raw uint8 batches. Augmentation runs on the GPU (see below) because\n    Kaggle only has ~4 CPU cores, which would otherwise bottleneck PIL-based augmentation.\"\"\"\n    g = torch.Generator(); g.manual_seed(SEED)\n    train_loader = make_train_loader(x_train, y_train, g)\n    val_loader = None\n    if x_val is not None:\n        val_loader = DataLoader(TensorDataset(x_val, y_val), batch_size=500, shuffle=False,\n                                num_workers=0, pin_memory=USE_AMP)      # deterministic, no augmentation\n    return train_loader, val_loader\n\n\ndef make_train_loader(x, y, generator=None):\n    return DataLoader(TensorDataset(x, y), batch_size=BATCH_SIZE, shuffle=True, drop_last=True,\n                      num_workers=NUM_WORKERS, pin_memory=USE_AMP, generator=generator,\n                      persistent_workers=NUM_WORKERS > 0)\n\n\ndata = get_datasets()\nclasses = data[\"classes\"]\ntrain_loader, val_loader = get_dataloaders(data[\"x_train\"], data[\"y_train\"], data[\"x_val\"], data[\"y_val\"])\nprint(\"Classes:\", classes)\nprint(f\"Train images: {len(data['x_train'])} | Validation images: {len(data['x_val'])}\")\nprint(\"Train class counts:\", torch.bincount(data[\"y_train\"]).tolist())\nprint(\"Val   class counts:\", torch.bincount(data[\"y_val\"]).tolist())\nprint(\"mean:\", data[\"mean\"].flatten().tolist(), \"std:\", data[\"std\"].flatten().tolist())\n\n# %% [markdown]\n# ## 4. GPU augmentation (training only) and deterministic preprocessing (validation / test)\n# * **Random crop (pad 4, reflect) + horizontal flip** - the standard CIFAR baseline; teaches translation/mirror invariance.\n# * **Cutout / random erasing** - hides a random patch so the network cannot rely on a single local cue.\n# * **CutMix** - pastes a patch from another image and mixes the labels; a strong regulariser for ResNets on CIFAR.\n# * **Label smoothing (0.1)** - reduces over-confidence.\n#\n# Validation/test data only get: uint8 -> float, /255, mean/std normalisation.\n\n# %%\ndef preprocess(xb_uint8, mean, std):\n    \"\"\"Deterministic preprocessing: NHWC uint8 -> normalised NCHW float (used for val/test).\"\"\"\n    x = xb_uint8.permute(0, 3, 1, 2).float().div_(255)\n    return ((x - mean) / std).contiguous(memory_format=torch.channels_last)\n\n\ndef random_crop(x, pad=4):\n    B, C, H, W = x.shape\n    xp = F.pad(x, (pad,) * 4, mode=\"reflect\")\n    i = torch.randint(0, 2 * pad + 1, (B,), device=x.device)\n    j = torch.randint(0, 2 * pad + 1, (B,), device=x.device)\n    rows = i[:, None] + torch.arange(H, device=x.device)            # (B,H)\n    cols = j[:, None] + torch.arange(W, device=x.device)            # (B,W)\n    b = torch.arange(B, device=x.device)[:, None, None, None]\n    c = torch.arange(C, device=x.device)[None, :, None, None]\n    return xp[b, c, rows[:, None, :, None], cols[:, None, None, :]]\n\n\ndef random_flip(x):\n    m = torch.rand(x.shape[0], device=x.device) < 0.5\n    return torch.where(m[:, None, None, None], x.flip(3), x)\n\n\ndef cutout(x, size):\n    \"\"\"Zero (= dataset mean after normalisation) a size x size square at a random centre, per image.\"\"\"\n    B, _, H, W = x.shape\n    cy = torch.randint(0, H, (B, 1, 1), device=x.device)\n    cx = torch.randint(0, W, (B, 1, 1), device=x.device)\n    ys = torch.arange(H, device=x.device).view(1, H, 1)\n    xs = torch.arange(W, device=x.device).view(1, 1, W)\n    m = (ys >= cy - size // 2) & (ys < cy + size - size // 2) & (xs >= cx - size // 2) & (xs < cx + size - size // 2)\n    return x.masked_fill(m[:, None], 0.0)\n\n\ndef gpu_augment(xb_uint8, aug, mean, std):\n    x = xb_uint8.permute(0, 3, 1, 2).float().div_(255)\n    x = (x - mean) / std\n    x = random_flip(random_crop(x))\n    if aug.get(\"cutout\", 0) > 0:\n        x = cutout(x, aug[\"cutout\"])\n    return x.contiguous(memory_format=torch.channels_last)\n\n\ndef mix_batch(x, y, aug):\n    \"\"\"CutMix / MixUp on a whole batch. Returns (x, y_a, y_b, lam); loss = lam*CE(y_a)+(1-lam)*CE(y_b).\"\"\"\n    mode = aug.get(\"mix\", \"none\")\n    if mode == \"none\" or random.random() > aug.get(\"mix_prob\", 0.5):\n        return x, y, y, 1.0\n    perm = torch.randperm(x.shape[0], device=x.device)\n    if mode == \"cutmix\":\n        lam = np.random.beta(1.0, 1.0)\n        H, W = x.shape[2:]\n        cut = math.sqrt(1.0 - lam)\n        ch, cw = int(H * cut), int(W * cut)\n        cy, cx = np.random.randint(H), np.random.randint(W)\n        y1, y2 = max(cy - ch // 2, 0), min(cy + ch // 2, H)\n        x1, x2 = max(cx - cw // 2, 0), min(cx + cw // 2, W)\n        x[:, :, y1:y2, x1:x2] = x[perm][:, :, y1:y2, x1:x2]\n        lam = 1.0 - (y2 - y1) * (x2 - x1) / (H * W)                  # exact kept-area ratio\n    elif mode == \"mixup\":\n        lam = np.random.beta(0.4, 0.4)\n        x = lam * x + (1.0 - lam) * x[perm]\n    else:\n        raise ValueError(mode)\n    return x, y, y[perm], lam\n\n# %% [markdown]\n# ## 5. Model: ResNet-18-style ConvNet adapted to 32x32 inputs, trained from scratch\n# 3x3 stem (no max-pool), 4 stages x 2 residual blocks (widths 64-128-256-512), BatchNorm, ReLU,\n# global average pooling, linear classifier. Kaiming-normal random init; the last BN in each residual block is\n# zero-initialised so blocks start as identity (stabilises early training). **No weights are ever loaded from outside.**\n\n# %%\nclass BasicBlock(nn.Module):\n    def __init__(self, in_ch, out_ch, stride):\n        super().__init__()\n        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, stride, 1, bias=False)\n        self.bn1 = nn.BatchNorm2d(out_ch)\n        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, 1, 1, bias=False)\n        self.bn2 = nn.BatchNorm2d(out_ch)\n        self.shortcut = nn.Sequential()\n        if stride != 1 or in_ch != out_ch:\n            self.shortcut = nn.Sequential(nn.Conv2d(in_ch, out_ch, 1, stride, bias=False), nn.BatchNorm2d(out_ch))\n\n    def forward(self, x):\n        out = F.relu(self.bn1(self.conv1(x)), inplace=True)\n        out = self.bn2(self.conv2(out))\n        return F.relu(out + self.shortcut(x), inplace=True)\n\n\nclass ResNetCIFAR(nn.Module):\n    def __init__(self, blocks=(2, 2, 2, 2), widths=(64, 128, 256, 512), num_classes=10, dropout=0.0):\n        super().__init__()\n        self.stem = nn.Sequential(nn.Conv2d(3, widths[0], 3, 1, 1, bias=False), nn.BatchNorm2d(widths[0]), nn.ReLU(inplace=True))\n        layers, in_ch = [], widths[0]\n        for stage, (n, w) in enumerate(zip(blocks, widths)):\n            for k in range(n):\n                layers.append(BasicBlock(in_ch, w, stride=2 if (k == 0 and stage > 0) else 1))\n                in_ch = w\n        self.layers = nn.Sequential(*layers)\n        self.pool = nn.AdaptiveAvgPool2d(1)\n        self.drop = nn.Dropout(dropout)\n        self.fc = nn.Linear(in_ch, num_classes)\n        for m in self.modules():                                     # random initialisation\n            if isinstance(m, nn.Conv2d):\n                nn.init.kaiming_normal_(m.weight, mode=\"fan_out\", nonlinearity=\"relu\")\n            elif isinstance(m, nn.BatchNorm2d):\n                nn.init.ones_(m.weight); nn.init.zeros_(m.bias)\n        for m in self.modules():\n            if isinstance(m, BasicBlock):\n                nn.init.zeros_(m.bn2.weight)\n\n    def forward(self, x):\n        x = self.layers(self.stem(x))\n        return self.fc(self.drop(torch.flatten(self.pool(x), 1)))\n\n\ndef build_model():\n    \"\"\"Custom ConvNet, randomly initialised. (If a torchvision model were used it would be\n    created with weights=None; here we do not use torchvision models at all.)\"\"\"\n    return ResNetCIFAR().to(device).to(memory_format=torch.channels_last)\n\n\ndef count_params(model):\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\n_m = build_model()\nprint(_m)\nprint(f\"Trainable parameters: {count_params(_m):,}\")\ndel _m\n\n# %% [markdown]\n# ## 6. Training / validation functions\n\n# %%\ndef make_scheduler(optimizer, total_steps, warmup_steps):\n    \"\"\"Per-iteration linear warm-up, then cosine decay to 0.\"\"\"\n    def f(step):\n        if step < warmup_steps:\n            return (step + 1) / warmup_steps\n        p = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n        return 0.5 * (1 + math.cos(math.pi * p))\n    return LambdaLR(optimizer, f)\n\n\ndef forward_tta(model, x, tta):\n    logits = model(x).float()\n    if tta:                                   # horizontal-flip test-time augmentation\n        logits = (logits + model(x.flip(3)).float()) / 2\n    return logits\n\n\ndef train_one_epoch(model, loader, optimizer, scheduler, scaler, aug):\n    model.train()\n    criterion = nn.CrossEntropyLoss(label_smoothing=aug.get(\"ls\", 0.0))\n    loss_sum = torch.zeros((), device=device); correct = torch.zeros((), device=device); n = 0\n    for xb, yb in loader:\n        xb = xb.to(device, non_blocking=True); yb = yb.to(device, non_blocking=True)\n        x = gpu_augment(xb, aug, data[\"mean\"], data[\"std\"])\n        x, ya, yb2, lam = mix_batch(x, yb, aug)\n        with torch.autocast(device_type=device.type, enabled=USE_AMP):\n            out = model(x)\n            loss = lam * criterion(out, ya) + (1 - lam) * criterion(out, yb2)\n        optimizer.zero_grad(set_to_none=True)\n        scaler.scale(loss).backward()\n        scaler.step(optimizer); scaler.update(); scheduler.step()\n        pred = out.argmax(1)\n        loss_sum += loss.detach() * len(yb)\n        correct += (pred == ya).sum() * lam + (pred == yb2).sum() * (1 - lam)   # accuracy w.r.t. mixed labels\n        n += len(yb)\n    return loss_sum.item() / n, correct.item() / n\n\n\n@torch.no_grad()\ndef validate(model, loader, tta=False):\n    \"\"\"Evaluation on the validation split: no augmentation, eval-mode BN, no gradients.\"\"\"\n    model.eval()\n    loss_sum, correct, n = 0.0, 0, 0\n    for xb, yb in loader:\n        x = preprocess(xb.to(device, non_blocking=True), data[\"mean\"], data[\"std\"])\n        yb = yb.to(device, non_blocking=True)\n        with torch.autocast(device_type=device.type, enabled=USE_AMP):\n            logits = forward_tta(model, x, tta)\n        loss_sum += F.cross_entropy(logits, yb, reduction=\"sum\").item()\n        correct += (logits.argmax(1) == yb).sum().item()\n        n += len(yb)\n    return loss_sum / n, correct / n\n\n\ndef train_model(train_loader, val_loader, aug, epochs, tag, save_ckpt=True, verbose=True):\n    \"\"\"Train from random init. If val_loader is given, validate every epoch and checkpoint the best\n    epoch by validation accuracy. If val_loader is None (full-data retrain) the final weights are saved.\"\"\"\n    model = build_model()\n    optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE, momentum=MOMENTUM,\n                                weight_decay=WEIGHT_DECAY, nesterov=True)\n    steps_per_epoch = len(train_loader)\n    scheduler = make_scheduler(optimizer, epochs * steps_per_epoch, WARMUP_EPOCHS * steps_per_epoch)\n    scaler = torch.amp.GradScaler(\"cuda\", enabled=USE_AMP)\n    best_acc, best_epoch, history = -1.0, 0, []\n    ckpt = os.path.join(WORK_DIR, f\"{tag}_best.pt\")\n    t0 = time.time()\n    for epoch in range(1, epochs + 1):\n        lr = optimizer.param_groups[0][\"lr\"]\n        tr_loss, tr_acc = train_one_epoch(model, train_loader, optimizer, scheduler, scaler, aug)\n        row = dict(epoch=epoch, lr=lr, train_loss=tr_loss, train_acc=tr_acc)\n        if val_loader is not None:\n            v_loss, v_acc = validate(model, val_loader)\n            row.update(val_loss=v_loss, val_acc=v_acc)\n            if v_acc > best_acc:\n                best_acc, best_epoch = v_acc, epoch\n                if save_ckpt: torch.save(model.state_dict(), ckpt)\n        history.append(row)\n        if verbose or epoch % 5 == 0 or epoch == epochs:\n            msg = f\"[{tag}] ep {epoch:3d}/{epochs} lr {lr:.4f} | train loss {tr_loss:.3f} acc {tr_acc*100:5.2f}%\"\n            if val_loader is not None:\n                msg += f\" | val loss {row['val_loss']:.3f} acc {row['val_acc']*100:5.2f}% (best {best_acc*100:.2f}% @ {best_epoch})\"\n            print(msg + f\" | {time.time()-t0:.0f}s\")\n    if val_loader is None and save_ckpt:\n        torch.save(model.state_dict(), os.path.join(WORK_DIR, f\"{tag}_final.pt\"))\n    hist = pd.DataFrame(history)\n    if save_ckpt: hist.to_csv(os.path.join(WORK_DIR, f\"{tag}_history.csv\"), index=False)\n    return model, hist, best_acc, best_epoch\n\n\ndef load_best_model(path):\n    \"\"\"Load a checkpoint produced by THIS notebook (our own trained weights, not pretrained ones).\"\"\"\n    model = build_model()\n    model.load_state_dict(torch.load(path, map_location=device, weights_only=True))\n    return model.eval()\n\n# %% [markdown]\n# ## 7. Small experiment study (validation set only)\n# Short runs comparing augmentation / regularisation settings with the same architecture, optimiser and seed.\n# The configuration with the best validation accuracy is used for the final run.\n# (Short runs may under-rate CutMix/label smoothing, which usually pay off with longer training - this is a limitation of the study.)\n\n# %%\nEXPERIMENTS = {\n    \"A_baseline_crop_flip\": dict(cutout=0, mix=\"none\",   ls=0.0),\n    \"B_cutout_ls\":          dict(cutout=8, mix=\"none\",   ls=0.1),\n    \"C_cutmix_ls\":          dict(cutout=0, mix=\"cutmix\", ls=0.1, mix_prob=0.5),\n}\n\nexp_results = {}\nif RUN_EXPERIMENTS:\n    for name, aug in EXPERIMENTS.items():\n        set_seed(SEED)\n        _, _, acc, ep = train_model(train_loader, val_loader, aug, EXP_EPOCHS, f\"exp_{name}\", save_ckpt=False, verbose=False)\n        exp_results[name] = (acc, ep)\n        torch.cuda.empty_cache()\n    print(\"\\nExperiment results (validation accuracy):\")\n    print(pd.DataFrame(exp_results, index=[\"best_val_acc\", \"best_epoch\"]).T.assign(\n        best_val_acc=lambda d: (d.best_val_acc * 100).round(2)).to_string())\n    FINAL_AUG_NAME = max(exp_results, key=lambda k: exp_results[k][0])\nelse:\n    FINAL_AUG_NAME = DEFAULT_AUG\nFINAL_AUG = EXPERIMENTS[FINAL_AUG_NAME]\nprint(\"\\nSelected augmentation config:\", FINAL_AUG_NAME, FINAL_AUG)\n\n# %% [markdown]\n# ## 8. Final training on the 80% split (validate every epoch, keep best epoch)\n\n# %%\nset_seed(SEED)\nmodel, history, best_val_acc, best_epoch = train_model(train_loader, val_loader, FINAL_AUG, EPOCHS, \"final_split\")\nlast_val_acc = history[\"val_acc\"].iloc[-1]\nprint(f\"\\nBest validation accuracy: {best_val_acc*100:.2f}%\")\nprint(f\"Best epoch: {best_epoch}\")\nprint(f\"Last-epoch validation accuracy: {last_val_acc*100:.2f}%\")\n\n# %%\nimport matplotlib.pyplot as plt\nfig, ax = plt.subplots(1, 3, figsize=(15, 4))\nax[0].plot(history.epoch, history.train_loss, label=\"train\"); ax[0].plot(history.epoch, history.val_loss, label=\"val\"); ax[0].set_title(\"Loss\"); ax[0].legend()\nax[1].plot(history.epoch, history.train_acc, label=\"train (vs mixed labels)\"); ax[1].plot(history.epoch, history.val_acc, label=\"val\"); ax[1].set_title(\"Accuracy\"); ax[1].legend()\nax[2].plot(history.epoch, history.lr); ax[2].set_title(\"Learning rate\")\nfor a in ax: a.set_xlabel(\"epoch\")\nplt.tight_layout(); plt.savefig(os.path.join(WORK_DIR, \"training_curves.png\"), dpi=120); plt.show()\n\n# %% [markdown]\n# ## 9. Final validation evaluation of the best checkpoint (+ measure whether flip-TTA helps)\n\n# %%\nbest_model = load_best_model(os.path.join(WORK_DIR, \"final_split_best.pt\"))\nval_loss_plain, val_acc_plain = validate(best_model, val_loader, tta=False)\nval_loss_tta,   val_acc_tta   = validate(best_model, val_loader, tta=True)\nprint(f\"Best checkpoint, no TTA   : val acc {val_acc_plain*100:.2f}% (loss {val_loss_plain:.3f})\")\nprint(f\"Best checkpoint, flip TTA : val acc {val_acc_tta*100:.2f}% (loss {val_loss_tta:.3f})\")\nUSE_TTA = val_acc_tta > val_acc_plain          # only use TTA if it measurably helps on validation\nprint(\"Use TTA for test predictions:\", USE_TTA)\n\n# %% [markdown]\n# ## 10. Optional: retrain the *selected* configuration on all 50,000 training images\n# Architecture, augmentation, epochs and all hyper-parameters stay fixed - nothing is re-tuned. There is no held-out data in this\n# stage, so the final-epoch weights are used (the cosine schedule ends at LR~0, so the last epoch is normally close to the best;\n# compare `Last-epoch validation accuracy` above). Official CIFAR-10 *test* labels are never used.\n# Accuracy of this model on the (now training) validation images is meaningless and is **not** reported.\n\n# %%\nfull_model = None\nif RETRAIN_FULL:\n    set_seed(SEED)\n    full_loader, _ = get_dataloaders(data[\"x_all\"], data[\"y_all\"], None, None)\n    print(f\"Retraining on {len(data['x_all'])} images (100% of official training set)\")\n    full_model, _, _, _ = train_model(full_loader, None, FINAL_AUG, EPOCHS, \"final_full\")\n    full_model.eval()\n\n# %% [markdown]\n# ## 11. Kaggle test set: load, predict, create submission\n\n# %%\nsample = pd.read_csv(SAMPLE_CSV)\nprint(\"Sample submission columns:\", list(sample.columns), \"| rows:\", len(sample))\nprint(sample.head(3))\n\ntest_dir = prepare_image_dir(\"test\", len(sample))\nn_png = sum(f.endswith(\".png\") for f in os.listdir(test_dir))\nprint(f\"PNG files in test dir: {n_png} | expected (rows in sample submission): {len(sample)}\")\nassert n_png == len(sample), \"test image count does not match sampleSubmission.csv\"\n\ntest_ids = sample[sample.columns[0]].values          # keep the exact id order/format of the sample file\nt0 = time.time(); x_test = load_images(test_dir, test_ids); print(f\"Loaded test images {tuple(x_test.shape)} in {time.time()-t0:.0f}s\")\n\n\n@torch.no_grad()\ndef predict_kaggle_test(model, x, tta, bs=1000):\n    model.eval(); preds = []\n    for i in range(0, len(x), bs):\n        xb = preprocess(x[i:i + bs].to(device, non_blocking=True), data[\"mean\"], data[\"std\"])\n        with torch.autocast(device_type=device.type, enabled=USE_AMP):\n            logits = forward_tta(model, xb, tta)\n        preds.append(logits.softmax(1).cpu())\n    return torch.cat(preds)\n\n\ndef create_submission(probs, path):\n    pred_labels = [classes[i] for i in probs.argmax(1).tolist()]\n    sub = pd.DataFrame({sample.columns[0]: test_ids, sample.columns[1]: pred_labels})\n    assert list(sub.columns) == list(sample.columns) and sub.shape == sample.shape\n    assert (sub[sample.columns[0]].values == sample[sample.columns[0]].values).all()\n    assert set(sub[sample.columns[1]]) <= set(classes)\n    sub.to_csv(path, index=False)\n    return sub\n\n\nprobs_split = predict_kaggle_test(best_model, x_test, USE_TTA)\nsub_split = create_submission(probs_split, os.path.join(WORK_DIR, \"submission_split_model.csv\"))   # backup: 80/20 best checkpoint\n\nif full_model is not None:\n    probs_final = predict_kaggle_test(full_model, x_test, USE_TTA)\n    agreement = (probs_final.argmax(1) == probs_split.argmax(1)).float().mean().item()\n    print(f\"Sanity check - agreement between 80/20 model and full-data model on test predictions: {agreement*100:.2f}%\")\n    FINAL_MODEL_USED = \"full 50k retrain (final epoch)\"\nelse:\n    probs_final, FINAL_MODEL_USED = probs_split, \"80/20 best-validation checkpoint\"\n\nsubmission = create_submission(probs_final, SUBMISSION_PATH)\nprint(\"Final submission built from:\", FINAL_MODEL_USED, \"| TTA:\", USE_TTA)\ndisplay(submission.head())\nprint(\"Shape:\", submission.shape)\nprint(submission[sample.columns[1]].value_counts())\nprint(\"Saved:\", SUBMISSION_PATH)\n\n# %% [markdown]\n# ## 12. Assignment compliance verification and report numbers\n\n# %%\nval_train_overlap = len(set(data[\"idx_train\"].tolist()) & set(data[\"idx_val\"].tolist()))\nmodel_for_report = build_model()\nprint(\"=\" * 60)\nprint(\"ConvNet:                              YES (ResNet-18-style CNN)\")\nprint(\"Weights randomly initialized:         YES (Kaiming init; custom model, nothing loaded before training)\")\nprint(\"Pretrained weights:                   NO\")\nprint(\"Transfer learning:                    NO\")\nprint(f\"80/20 train-validation split:         YES ({len(data['x_train'])} / {len(data['x_val'])})\")\nprint(f\"Validation images used for training:  {'NO' if val_train_overlap == 0 else 'YES (BUG!)'} (index overlap in 80/20 runs = {val_train_overlap})\")\nprint(f\"Kaggle submission generated:          {'YES' if os.path.exists(SUBMISSION_PATH) and submission.shape == sample.shape else 'NO'}\")\nprint(\"Note: the optional full-data retrain (if enabled) uses all 50k OFFICIAL TRAIN images after model selection; test images never used for training.\")\nprint(\"=\" * 60)\nprint(model_for_report)\nprint(f\"Total trainable parameters : {count_params(model_for_report):,}\")\nprint(f\"Training images            : {len(data['x_train'])}\")\nprint(f\"Validation images          : {len(data['x_val'])}\")\nprint(f\"Best validation accuracy   : {best_val_acc*100:.2f}%   (with flip-TTA: {val_acc_tta*100:.2f}%)\")\nprint(f\"Best epoch                 : {best_epoch}\")\nprint(f\"Last-epoch val accuracy    : {last_val_acc*100:.2f}%\")\nprint(\"Experiment results         :\", {k: f\"{v[0]*100:.2f}% @ep{v[1]}\" for k, v in exp_results.items()} or \"not run\")\nprint(\"Final configuration        :\", dict(seed=SEED, batch_size=BATCH_SIZE, epochs=EPOCHS, optimizer=\"SGD nesterov\",\n                                            momentum=MOMENTUM, peak_lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY,\n                                            warmup_epochs=WARMUP_EPOCHS, schedule=\"warmup+cosine\", aug=FINAL_AUG_NAME, aug_params=FINAL_AUG,\n                                            mixed_precision=USE_AMP, tta=USE_TTA, final_model_used=FINAL_MODEL_USED))\nprint(\"Kaggle public/private score: <fill in manually after submitting - not computed by this notebook>\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-09-19T09:43:43.239361Z","iopub.execute_input":"2026-09-19T09:43:43.239847Z","iopub.status.idle":"2026-09-19T14:13:58.509634Z","shell.execute_reply.started":"2026-09-19T09:43:43.23981Z","shell.execute_reply":"2026-09-19T14:13:58.508847Z"}},"outputs":[],"execution_count":null}]}