{"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# * Optimizer: **AdamW** (decoupled weight decay)                                        # MODIFIED\n# * LR schedule: **linear warm-up -> cosine annealing** (per iteration)                  # MODIFIED\n# * Augmentation: **RandomCrop + RandomHorizontalFlip + CutMix**                         # MODIFIED\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# * The **best-validation-accuracy checkpoint** is loaded and used for the Kaggle test predictions   # MODIFIED\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, io, sys, time, glob, shutil, random, subprocess, math   # MODIFIED: + io\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        = 300         # MODIFIED: 60 -> 200; strong augmentation (CutMix + TrivialAugment) needs long training.\n                            #   This is the single biggest accuracy gain. Use 100 if GPU time is short (~-0.5%).\nWIDTHS        = (96, 192, 384, 768)   # MODIFIED: channel widths of the 4 stages (ResNet-18 uses 64-128-256-512).\n                            #   1.5x wider -> ~25M params, more capacity for 200 epochs of heavy augmentation.\nUSE_TRIVIALAUG = True       # MODIFIED: TrivialAugmentWide (one random photometric/geometric op per image) on top\n                            #   of crop + flip + CutMix; the standard strong augmentation policy for CIFAR\nLEARNING_RATE = 2e-3        # MODIFIED: peak LR for AdamW (was 0.1 for SGD)\nWEIGHT_DECAY  = 5e-2        # MODIFIED: AdamW decoupled weight decay. Per step AdamW shrinks weights by lr*wd,\n                            #   so it needs much larger values than SGD (1e-4 here gave ~no regularisation -> 88%)\nADAM_BETAS    = (0.9, 0.999)  # MODIFIED: AdamW betas (replaces SGD MOMENTUM, which is removed)\nWARMUP_EPOCHS = 2           # linear warm-up length, then cosine annealing to 0\nNUM_WORKERS   = min(4, os.cpu_count() or 2)   # MODIFIED: TrivialAugment runs in the loader workers (CPU)\nVAL_FRACTION  = 0.2         # -> 40,000 train / 10,000 validation\n\nRUN_EXPERIMENTS = False     # MODIFIED: optional ablation study (reported only, see section 7); off to save ~17 min GPU\nEXP_EPOCHS      = 15        # epochs per experiment (short runs to save GPU time)\nFINAL_AUG_NAME  = \"C_cutmix_ls\"   # MODIFIED: final run always uses RandomCrop + Flip + CutMix (assignment requirement)\nRETRAIN_FULL    = False     # MODIFIED: off by default; the submission must come from the best-validation checkpoint\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\")\nBEST_CKPT_PATH  = os.path.join(WORK_DIR, \"final_split_best.pt\")   # MODIFIED: single name for the best checkpoint\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 / load the Kaggle data\n# The competition provides `train.7z` (50k PNGs), `test.7z` (300k PNGs), `trainLabels.csv`, `sampleSubmission.csv`.\n# MODIFIED: images are **streamed straight from the .7z archive into memory** (`7z x -so`), so no\n# 300k tiny files are written to disk (that took 30-40 min on Kaggle). If image folders already exist they are used;\n# disk extraction is kept only as a fallback when no 7z binary is available.\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_from_7z(archive, ids, exe):                          # MODIFIED: new loader\n    \"\"\"Stream PNGs out of a .7z archive and decode them in memory - nothing is written to disk.\n    `7z l -slt` gives every file's name and byte size in archive order; `7z x -so` writes all file\n    contents back-to-back to stdout in that same order, so we cut the stream into files by size.\"\"\"\n    listing = subprocess.run([exe, \"l\", \"-slt\", archive], check=True, capture_output=True, text=True).stdout\n    entries, path = [], None\n    for line in listing.splitlines():\n        if line.startswith(\"Path = \"):\n            path = line[7:]\n        elif line.startswith(\"Size = \") and path is not None:\n            entries.append((path, int(line[7:]))); path = None\n    entries = [(p, s) for p, s in entries if p.lower().endswith(\".png\")]\n    pos = {str(i): k for k, i in enumerate(ids)}                     # image id -> row in the output tensor\n    out = np.empty((len(ids), 32, 32, 3), np.uint8)\n\n    def _decode(k, buf):\n        out[k] = np.asarray(Image.open(io.BytesIO(buf)).convert(\"RGB\"))\n\n    proc = subprocess.Popen([exe, \"x\", \"-so\", archive], stdout=subprocess.PIPE,\n                            stderr=subprocess.DEVNULL, bufsize=1 << 20)\n    seen, t0 = 0, time.time()\n    with ThreadPoolExecutor(max_workers=8) as ex:\n        futures = []\n        for p, size in entries:\n            buf = proc.stdout.read(size)\n            assert len(buf) == size, f\"archive stream ended early at {p}\"\n            k = pos.get(os.path.splitext(os.path.basename(p))[0])\n            if k is not None:\n                futures.append(ex.submit(_decode, k, buf)); seen += 1\n                if seen % 50000 == 0:\n                    print(f\"  streamed {seen}/{len(ids)} images, {time.time()-t0:.0f}s\", flush=True)\n        for f in futures:\n            f.result()\n    assert proc.wait() == 0, \"7z failed while streaming\"\n    assert seen == len(ids), f\"found {seen} of {len(ids)} images in {archive}\"\n    return torch.from_numpy(out)\n\n\ndef load_split(split, ids):                                          # MODIFIED: new entry point\n    \"\"\"Return a uint8 tensor (N, 32, 32, 3) for split in {'train','test'}, rows in the order of `ids`.\n    Uses already-extracted folders if present, otherwise streams from the .7z (fast path).\"\"\"\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) >= len(ids):\n            return load_images(d, ids)\n    archive = os.path.join(COMP_DIR, f\"{split}.7z\")\n    exe = shutil.which(\"7z\") or shutil.which(\"7zz\") or shutil.which(\"7za\") or shutil.which(\"7zr\")\n    if exe is not None and os.path.exists(archive):\n        print(f\"Streaming {archive} into memory with {exe} ...\")\n        t0 = time.time(); x = load_images_from_7z(archive, ids, exe)\n        print(f\"  {len(x)} images loaded in {time.time()-t0:.0f}s\")\n        return x\n    return load_images(prepare_image_dir(split, len(ids)), ids)       # fallback: extract to disk\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    x_all = load_split(\"train\", labels[\"id\"].values)                # MODIFIED: stream from archive\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\nclass TrivialAugDataset(torch.utils.data.Dataset):                   # MODIFIED: new class\n    \"\"\"Training images with TrivialAugmentWide applied per image (in the DataLoader workers).\n    Each image gets ONE randomly chosen op (rotate, shear, translate, brightness, contrast, colour,\n    sharpness, posterize, solarize, autocontrast, equalize or identity) with a random strength.\"\"\"\n    def __init__(self, x, y):\n        from torchvision.transforms import TrivialAugmentWide\n        self.x, self.y, self.tf = x, y, TrivialAugmentWide()\n\n    def __len__(self):\n        return len(self.x)\n\n    def __getitem__(self, i):\n        img = self.tf(self.x[i].permute(2, 0, 1)).permute(1, 2, 0)   # HWC uint8 -> CHW for torchvision -> back\n        return img, self.y[i]\n\n\ndef make_train_loader(x, y, generator=None):\n    ds = TrivialAugDataset(x, y) if USE_TRIVIALAUG else TensorDataset(x, y)   # MODIFIED\n    return DataLoader(ds, 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# * **CutMix** (advanced augmentation used in the final model) - pastes a random patch from another image of the batch\n#   and mixes the two labels in proportion to the patch area; a strong regulariser for ResNets on CIFAR.   # MODIFIED\n# * Cutout / random erasing - only used in the optional ablation (config B), not in the final model.       # MODIFIED\n# * Label smoothing (0.1) - reduces over-confidence.\n# * TrivialAugmentWide (CPU, in the loader workers, before crop/flip) - one random colour/geometry op per image.  # MODIFIED\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    \"\"\"RandomCrop(32, padding=4, reflect), independently per image, on the GPU.\"\"\"\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    \"\"\"RandomHorizontalFlip(p=0.5), independently per image, on the GPU.\"\"\"\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    CutMix: lam ~ Beta(1,1); a box covering ~(1-lam) of the image is copied from a shuffled copy of the\n    batch, and lam is then recomputed as the exact fraction of the original image that is kept.\"\"\"\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: custom wide residual ConvNet for 32x32 inputs, trained from scratch                    # MODIFIED\n# 3x3 stem (no max-pool), 4 stages x 2 residual blocks (widths 96-192-384-768, set by WIDTHS), 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(widths=WIDTHS).to(device).to(memory_format=torch.channels_last)   # MODIFIED: custom widths\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# **Optimizer:** AdamW (Adam with decoupled weight decay).                                   # MODIFIED\n# **LR schedule (per iteration):**                                                            # MODIFIED\n# * warm-up, steps `0 .. W-1`:  `lr = peak_lr * (step + 1) / W`  (linear increase to the peak)\n# * cosine,  steps `W .. T`:    `lr = peak_lr * 0.5 * (1 + cos(pi * (step - W) / (T - W)))`  (decay to 0)\n\n# %%\ndef make_scheduler(optimizer, total_steps, warmup_steps):\n    \"\"\"MODIFIED (documented): linear warm-up followed by cosine annealing, updated every iteration.\n    LambdaLR multiplies the optimizer's peak LR by the factor returned here.\"\"\"\n    def lr_factor(step):\n        # Phase 1 - linear warm-up: factor rises from 1/W to 1 over the first W steps\n        if step < warmup_steps:\n            return (step + 1) / warmup_steps\n        # Phase 2 - cosine annealing: factor falls from 1 to 0 over the remaining steps\n        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n        return 0.5 * (1 + math.cos(math.pi * min(1.0, progress)))\n    return LambdaLR(optimizer, lr_factor)\n\n\ndef make_optimizer(model):\n    \"\"\"MODIFIED: AdamW with decoupled weight decay (replaces SGD + Nesterov momentum).\n    Weight decay is applied to conv/linear weights only. BatchNorm scales/shifts and biases are excluded:\n    they don't cause overfitting, and decaying them at wd=0.05 (incl. the zero-initialised BN in each block)\n    would just fight the normalisation. This is the standard way AdamW is used.\"\"\"\n    decay    = [p for p in model.parameters() if p.ndim > 1]        # conv + linear weight tensors\n    no_decay = [p for p in model.parameters() if p.ndim <= 1]       # BN weights/biases, linear bias\n    return torch.optim.AdamW([{\"params\": decay, \"weight_decay\": WEIGHT_DECAY},\n                              {\"params\": no_decay, \"weight_decay\": 0.0}],\n                             lr=LEARNING_RATE, betas=ADAM_BETAS)\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\"])          # RandomCrop + RandomHorizontalFlip\n        x, ya, yb2, lam = mix_batch(x, yb, aug)                      # CutMix (with probability mix_prob)\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()    # LR schedule stepped every iteration\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()                                            # fresh random weights, nothing loaded\n    optimizer = make_optimizer(model)                                # MODIFIED: AdamW (no SGD anywhere)\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        saved = \"\"\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:                                     # validation accuracy improved\n                best_acc, best_epoch = v_acc, epoch\n                if save_ckpt:\n                    torch.save(model.state_dict(), ckpt)             # checkpoint on every improvement\n                    saved = \" | checkpoint saved\"                    # MODIFIED: make saving visible in the log\n        history.append(row)\n        if verbose or epoch % 5 == 0 or epoch == epochs:\n            msg = f\"[{tag}] ep {epoch:3d}/{epochs} lr {lr:.2e} | train loss {tr_loss:.3f} acc {tr_acc*100:5.2f}%\"   # MODIFIED: lr format for small AdamW LRs\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\" + saved)\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 ablation study (validation set only)                                           # MODIFIED\n# Short runs comparing augmentation settings with the same architecture, AdamW optimizer, schedule and seed.\n# These results are **reported only**: the final model always uses RandomCrop + Flip + CutMix (config C),\n# as required by the assignment. Set `RUN_EXPERIMENTS = False` to skip this section and save GPU time.\n# (Short runs may under-rate CutMix, which usually pays off with longer training.)\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),   # final configuration\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(\"\\nAblation 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# MODIFIED: no automatic selection any more - the final run always uses CutMix (FINAL_AUG_NAME set in section 1)\nFINAL_AUG = EXPERIMENTS[FINAL_AUG_NAME]\nassert FINAL_AUG.get(\"mix\") == \"cutmix\", \"final configuration must use CutMix\"\nprint(\"\\nFinal augmentation config:\", FINAL_AUG_NAME, FINAL_AUG)\n\n# %% [markdown]\n# ## 8. Final training on the 80% split (validate every epoch, save checkpoint whenever val accuracy improves)\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 (warm-up + cosine)\")   # MODIFIED: title\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. Load the best validation checkpoint and evaluate it on validation (+ measure whether flip-TTA helps)\n\n# %%\nbest_model = load_best_model(BEST_CKPT_PATH)                         # MODIFIED: shared checkpoint path constant\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 (epoch {best_epoch}), no TTA   : val acc {val_acc_plain*100:.2f}% (loss {val_loss_plain:.3f})\")\nprint(f\"Best checkpoint (epoch {best_epoch}), 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 (OFF by default): retrain the selected configuration on all 50,000 training images  # MODIFIED\n# Kept only for comparison. Its predictions are written to a **separate** file (`submission_full_retrain.csv`);\n# the main `submission.csv` always comes from the best-validation checkpoint, as required by the assignment.\n# Official CIFAR-10 *test* labels are never used.\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 the BEST VALIDATION CHECKPOINT, predict, create submission          # MODIFIED\n\n# %%\nsample = pd.read_csv(SAMPLE_CSV)\nprint(\"Sample submission columns:\", list(sample.columns), \"| rows:\", len(sample))\nprint(sample.head(3))\n\ntest_ids = sample[sample.columns[0]].values          # keep the exact id order/format of the sample file\nx_test = load_split(\"test\", test_ids)                # MODIFIED: streamed from test.7z; asserts every id is found\nassert len(x_test) == len(sample), \"test image count does not match sampleSubmission.csv\"\nprint(f\"Test images: {tuple(x_test.shape)}\")\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\n# MODIFIED: explicitly (re)load the best validation checkpoint right before test inference\nbest_model = load_best_model(BEST_CKPT_PATH)\nprint(f\"Loaded best validation checkpoint: {BEST_CKPT_PATH} (epoch {best_epoch}, val acc {best_val_acc*100:.2f}%)\")\nprobs_final = predict_kaggle_test(best_model, x_test, USE_TTA)\nFINAL_MODEL_USED = f\"80/20 best-validation checkpoint (epoch {best_epoch})\"\n\nif full_model is not None:                     # MODIFIED: optional comparison only, saved to a separate file\n    probs_full = predict_kaggle_test(full_model, x_test, USE_TTA)\n    create_submission(probs_full, os.path.join(WORK_DIR, \"submission_full_retrain.csv\"))\n    agreement = (probs_full.argmax(1) == probs_final.argmax(1)).float().mean().item()\n    print(f\"Agreement between best checkpoint and full-data retrain on test predictions: {agreement*100:.2f}%\")\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()\n_opt_check = make_optimizer(model_for_report)                        # MODIFIED: verify optimizer type\nprint(\"=\" * 60)\nprint(f\"ConvNet:                              YES (custom wide residual CNN, widths {WIDTHS})\")   # MODIFIED\nprint(\"Weights randomly initialized:         YES (Kaiming init; custom model, nothing loaded before training)\")\nprint(\"Pretrained weights:                   NO\")\nprint(\"Transfer learning:                    NO\")\nprint(f\"Optimizer:                            {type(_opt_check).__name__} (weight decay {WEIGHT_DECAY}, betas {ADAM_BETAS})\")   # MODIFIED\nprint(f\"SGD / SGD+momentum used:              {'YES (BUG!)' if isinstance(_opt_check, torch.optim.SGD) else 'NO'}\")             # MODIFIED\nprint(f\"LR schedule:                          linear warm-up ({WARMUP_EPOCHS} epochs) -> cosine annealing to 0, peak {LEARNING_RATE}\")  # MODIFIED\nprint(f\"Augmentation:                         {'TrivialAugmentWide + ' if USE_TRIVIALAUG else ''}RandomCrop(pad 4) + RandomHorizontalFlip + CutMix (p={FINAL_AUG.get('mix_prob')})\")       # MODIFIED\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 = {val_train_overlap})\")\nprint(\"Model selection:                      best validation accuracy (checkpoint saved on every improvement)\")                       # MODIFIED\nprint(f\"Test predictions from:                {FINAL_MODEL_USED}\")                                                                       # MODIFIED\nprint(f\"Kaggle submission generated:          {'YES' if os.path.exists(SUBMISSION_PATH) and submission.shape == sample.shape else 'NO'}\")\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(\"Ablation 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=\"AdamW\",   # MODIFIED\n                                            betas=ADAM_BETAS, peak_lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY,\n                                            warmup_epochs=WARMUP_EPOCHS, schedule=\"linear warmup + cosine annealing\",\n                                            aug=FINAL_AUG_NAME, aug_params=FINAL_AUG,\n                                            trivialaugment=USE_TRIVIALAUG, widths=WIDTHS,   # MODIFIED\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-10-02T11:19:00.292425Z","iopub.execute_input":"2026-10-02T11:19:00.292708Z","iopub.status.idle":"2026-10-02T13:27:08.795393Z","shell.execute_reply.started":"2026-10-02T11:19:00.292679Z","shell.execute_reply":"2026-10-02T13:27:08.794546Z"},"_kg_hide-input":false},"outputs":[],"execution_count":null}]}