{"cells":[{"cell_type":"markdown","id":"c00","metadata":{},"source":"# CIFAR-10 from scratch: PyraNet-ECA\n\nCompetition: <https://www.kaggle.com/c/cifar-10>\n\nThis notebook trains a convolutional net I put together myself on CIFAR-10,\nstarting from random weights. Nothing is pre-trained; the only thing fetched\nfrom anywhere is the image data itself.\n\nWhere each requirement of the assignment is handled:\n\n| Requirement | Where it lives |\n|---|---|\n| No pre-trained weights | §1 builds and initialises every layer by hand |\n| Advanced optimizer | AdamW, §3 |\n| Cosine schedule with a warmup phase | §3, stepped once per iteration |\n| Advanced augmentation | CutMix in §2, with MixUp and Cutout alongside it |\n| My own architecture, not off the shelf | PyraNet-ECA, §1 |\n\nThe design in one sentence: three stages at 32, 16 and 8 pixels, pre-activation\nresidual blocks with a very cheap form of channel attention, plus the handful of\nregularisers that only start to pay off once the schedule is long. This is the Kaggle version. It reads the competition archives straight out of\n`/kaggle/input`, trains from random weights, and writes\n`/kaggle/working/submission.csv`. No internet access needed.\n\nTwo things differ from the run I did offline. The schedule is 120 epochs\ninstead of 300, so it comfortably finishes inside one Kaggle GPU session. And\nsince the only labels Kaggle hands over are the 50 000 training ones, a\nstratified 5 000-image hold-out stands in for the accuracy curve."},{"cell_type":"code","execution_count":null,"id":"c01","metadata":{},"outputs":[],"source":"%matplotlib inline\nimport os, math, time, json, glob, shutil, subprocess\nfrom concurrent.futures import ThreadPoolExecutor\n\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\ntorch.backends.cudnn.benchmark = True\ntorch.backends.cuda.matmul.allow_tf32 = True\ntorch.backends.cudnn.allow_tf32 = True\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Pick the autocast dtype from the actual hardware. Careful here: on a Turing\n# card torch.cuda.is_bf16_supported() still returns True, but bf16 is emulated\n# and slow, so go by compute capability instead -- real bf16 starts at sm_80.\n# bf16 needs no loss scaling; fp16 does, hence the GradScaler further down.\n_native_bf16 = device.type == \"cuda\" and torch.cuda.get_device_capability(0) >= (8, 0)\nAMP_DTYPE = torch.bfloat16 if _native_bf16 else torch.float16\n\nprint(\"torch\", torch.__version__, \"|\", torch.cuda.get_device_name(0) if device.type == \"cuda\" else \"cpu\")\nprint(\"autocast dtype:\", AMP_DTYPE, \"| native bf16:\", _native_bf16)\n\n# Where things live on Kaggle. /kaggle/input is the read-only data mount,\n# /kaggle/working is the only directory kept as notebook output, and /tmp is\n# scratch with plenty of room. Don't hardcode the competition path: depending on\n# how the data is attached it shows up either as /kaggle/input/cifar-10 or as\n# /kaggle/input/competitions/cifar-10, so go and find it.\ndef find_competition_dir(root=\"/kaggle/input\"):\n    \"\"\"Locate the competition mount. Attaching the data can land it at\n    <root>/cifar-10 or <root>/competitions/cifar-10, so try both and then just\n    go looking for train.7z before giving up.\"\"\"\n    for cand in (os.path.join(root, \"competitions\", \"cifar-10\"),\n                 os.path.join(root, \"cifar-10\")):\n        if os.path.exists(os.path.join(cand, \"train.7z\")):\n            return cand\n    hits = sorted(glob.glob(os.path.join(root, \"**\", \"train.7z\"), recursive=True))\n    return os.path.dirname(hits[0]) if hits else os.path.join(root, \"cifar-10\")\n\nCOMP_DIR = os.environ.get(\"CIFAR10_COMP_DIR\") or find_competition_dir()\nSCRATCH  = os.environ.get(\"CIFAR10_SCRATCH\", \"/tmp\")\nOUT_DIR  = os.environ.get(\"CIFAR10_OUT\",\n                          \"/kaggle/working\" if os.path.isdir(\"/kaggle/working\") else \"./runs\")\nos.makedirs(SCRATCH, exist_ok=True)\nos.makedirs(OUT_DIR, exist_ok=True)\n\nif os.path.isdir(\"/kaggle/input\"):\n    print(\"under /kaggle/input:\", sorted(os.listdir(\"/kaggle/input\")))\nprint(\"competition:\", COMP_DIR,\n      \"->\", sorted(os.listdir(COMP_DIR)) if os.path.isdir(COMP_DIR) else \"MISSING\")\nprint(\"scratch    :\", SCRATCH)\nprint(\"output     :\", OUT_DIR)"},{"cell_type":"code","execution_count":null,"id":"c02","metadata":{},"outputs":[],"source":"CFG = dict(\n    # ---- architecture ----\n    arch           = \"base\",     # see CONFIGS in §1\n    # ---- optimisation ----\n    epochs         = 120,     # fits a Kaggle GPU session; see note above\n    val_n          = 5000,    # stratified hold-out for the accuracy curve\n    batch_size     = 256,\n    peak_lr        = 2e-3,       # AdamW peak LR at batch size 256\n    min_lr_frac    = 1e-3,       # cosine floor = peak_lr * this\n    weight_decay   = 0.05,       # decoupled, ≥2-D tensors only\n    betas          = (0.9, 0.99),\n    warmup_epochs  = 5,\n    label_smooth   = 0.1,\n    ema_decay      = 0.9998,\n    # ---- augmentation ----\n    pad            = 4,          # reflect-pad canvas for the random crop\n    affine_p       = 0.5,        # p(random rotate/shear/scale)\n    color_p        = 0.5,        # p(each colour op)\n    color_strength = 0.25,\n    cutout         = 12,         # Cutout square side, 0 disables\n    mix_prob       = 0.5,        # p(mix the batch at all)\n    cutmix_share   = 0.5,        # of mixed batches, fraction using CutMix vs MixUp\n    cutmix_alpha   = 1.0,\n    mixup_alpha    = 0.2,\n    # ---- runtime ----\n    compile_model  = True,\n    seed           = 0,\n)\n\n# Quick end-to-end pipeline check without editing the notebook:\n#     SMOKE=1 jupyter nbconvert --execute --inplace cifar10_pyranet.ipynb\nif os.environ.get(\"SMOKE\") == \"1\":\n    CFG.update(arch=\"small\", epochs=4, warmup_epochs=1, compile_model=False, val_n=1000)\n    print(\">>> SMOKE MODE: tiny run, results are not meaningful\")\n\nfor k, v in CFG.items():\n    print(f\"{k:>15} = {v}\")"},{"cell_type":"markdown","id":"c03","metadata":{},"source":"---\n## 1 · The network\n\nEverything here is written out below — no `torchvision.models`, no `timm`. The\nshape is a three-stage pyramid that stays at CIFAR's own resolutions:\n\n```\ninput 3x32x32\n  stem      three 3x3 convs, still 32x32\n  stage 1   D1 blocks at 32x32\n  stage 2   D2 blocks at 16x16   (first block strides)\n  stage 3   D3 blocks at  8x8    (first block strides)\n  BN, SiLU, global average pool, dropout, linear -> 10\n```\n\nSix things in here are deliberate choices rather than defaults, so it's worth\nsaying why I made them.\n\n**A deeper stem.** ImageNet models throw away resolution immediately; a 7x7\nstride-2 conv on a 32x32 image would be vandalism. I use three cheap 3x3 convs\nat full resolution instead of the single 3x3 you'd normally see on CIFAR. It\nbuys a richer set of low-level features for very few parameters.\n\n**Pre-activation blocks** (`BN -> SiLU -> conv` rather than conv-then-BN). The\npoint is that the skip connection stays a plain addition with nothing sitting on\nit, which is what lets the deeper variants train without me having to reach for\ngradient clipping.\n\n**ECA instead of SE for channel attention.** Squeeze-and-Excitation would cost\n`2C²/r` parameters per block. ECA treats the pooled channel vector as a 1-D\nsignal and slides one small 1-D convolution along it, so neighbouring channels\nstill get to interact but the cost is about `k` parameters — single digits —\nno matter how wide the layer is. Reweighting channels is worth doing on CIFAR\nbecause at 8x8 a lot of feature maps have already specialised to one class.\n\n**LayerScale on every residual branch,** a learnable per-channel vector starting\nat `1e-4`. At initialisation each block is therefore almost exactly the identity,\nso the network starts shallow in effect and deepens as the scales grow. This is\nthe thing that makes a 2e-3 AdamW learning rate survive the first few hundred\nsteps.\n\n**Stochastic depth that ramps with depth** rather than a flat rate: no drop at\nthe stem, the full rate at the last block. Late blocks are where the overfitting\nhappens, so they're the ones worth hitting hardest.\n\n**Downsampling without aliasing.** The usual stride-2 1x1 convolution on the\nshortcut path looks at one pixel in four and throws the rest away. Here the\nshortcut is `AvgPool(2)` followed by a 1x1, so it low-passes before it\nsubsamples. Reliably worth a little accuracy, costs nothing."},{"cell_type":"code","execution_count":null,"id":"c04","metadata":{},"outputs":[],"source":"def conv3x3(cin, cout, stride=1):\n    return nn.Conv2d(cin, cout, 3, stride=stride, padding=1, bias=False)\n\n\nclass ECA(nn.Module):\n    \"\"\"Channel attention on the cheap.\n\n    Squeeze-Excitation would need 2C^2/r parameters here. This pools to one\n    number per channel, then slides a single small 1-D conv along that vector,\n    so adjacent channels still talk to each other for about k parameters --\n    three to nine of them, whatever the width. k grows with log2(C).\n    \"\"\"\n\n    def __init__(self, channels, gamma=2.0, beta=1.0):\n        super().__init__()\n        k = int(abs(math.log2(channels) / gamma + beta / gamma))\n        k = max(3, k | 1)                       # force odd so the conv is centred\n        self.conv = nn.Conv1d(1, 1, k, padding=k // 2, bias=False)\n\n    def forward(self, x):\n        w = x.mean(dim=(2, 3))                  # B,C    global channel descriptor\n        w = self.conv(w.unsqueeze(1)).squeeze(1)  # B,C  local cross-channel mixing\n        return x * torch.sigmoid(w)[:, :, None, None]\n\n\nclass DropPath(nn.Module):\n    \"\"\"Stochastic depth: throw away the whole residual branch, per sample.\"\"\"\n\n    def __init__(self, p=0.0):\n        super().__init__()\n        self.p = float(p)\n\n    def forward(self, x):\n        if self.p == 0.0 or not self.training:\n            return x\n        keep = 1.0 - self.p\n        mask = x.new_empty(x.shape[0], 1, 1, 1).bernoulli_(keep)\n        return x * mask / keep                  # /keep keeps the expectation fixed\n\n    def extra_repr(self):\n        return f\"p={self.p:.3f}\"\n\n\nclass LayerScale(nn.Module):\n    \"\"\"Per-channel gain on the residual branch. Starts near zero, so the block\n    begins life as (almost) the identity and has to earn its contribution.\"\"\"\n\n    def __init__(self, channels, init=1e-4):\n        super().__init__()\n        self.gamma = nn.Parameter(torch.full((channels,), float(init)))\n\n    def forward(self, x):\n        return x * self.gamma[None, :, None, None]"},{"cell_type":"code","execution_count":null,"id":"c05","metadata":{},"outputs":[],"source":"class Block(nn.Module):\n    \"\"\"BN -> SiLU -> 3x3 -> BN -> SiLU -> 3x3 -> ECA -> LayerScale -> DropPath (+ shortcut).\"\"\"\n\n    def __init__(self, cin, cout, stride=1, drop_path=0.0, layer_scale=1e-4):\n        super().__init__()\n        self.bn1   = nn.BatchNorm2d(cin)\n        self.conv1 = conv3x3(cin, cout, stride)\n        self.bn2   = nn.BatchNorm2d(cout)\n        self.conv2 = conv3x3(cout, cout, 1)\n        self.eca   = ECA(cout)\n        self.scale = LayerScale(cout, layer_scale)\n        self.drop_path = DropPath(drop_path)\n\n        if stride != 1 or cin != cout:\n            # AvgPool *then* 1x1. A strided 1x1 would sample one pixel in four.\n            pool = nn.AvgPool2d(stride, stride) if stride != 1 else nn.Identity()\n            self.shortcut = nn.Sequential(pool, nn.Conv2d(cin, cout, 1, bias=False))\n        else:\n            self.shortcut = None                # clean identity path\n\n    def forward(self, x):\n        h = F.silu(self.bn1(x))\n        # ResNet-v2 convention: a projecting shortcut reads the pre-activated\n        # tensor, so both paths see the same normalised input.\n        identity = x if self.shortcut is None else self.shortcut(h)\n        h = self.conv1(h)\n        h = self.conv2(F.silu(self.bn2(h)))\n        h = self.scale(self.eca(h))\n        return identity + self.drop_path(h)"},{"cell_type":"code","execution_count":null,"id":"c06","metadata":{},"outputs":[],"source":"class PyraNet(nn.Module):\n    def __init__(self, num_classes=10, widths=(128, 256, 512), depths=(4, 4, 4),\n                 stem_width=64, drop_path=0.1, dropout=0.1, layer_scale=1e-4):\n        super().__init__()\n        assert len(widths) == len(depths) == 3\n\n        # --- deep stem, held at 32x32 ---\n        s = stem_width\n        self.stem = nn.Sequential(\n            conv3x3(3, s // 2),      nn.BatchNorm2d(s // 2), nn.SiLU(inplace=True),\n            conv3x3(s // 2, s // 2), nn.BatchNorm2d(s // 2), nn.SiLU(inplace=True),\n            conv3x3(s // 2, s),      # no BN/act: the first pre-act block applies them\n        )\n\n        # --- three stages ---\n        total, made, cin, stages = sum(depths), 0, s, []\n        for si, (w, d) in enumerate(zip(widths, depths)):\n            blocks = []\n            for b in range(d):\n                stride = 2 if (si > 0 and b == 0) else 1\n                dp = drop_path * made / max(1, total - 1)   # linear ramp with depth\n                blocks.append(Block(cin, w, stride, dp, layer_scale))\n                cin, made = w, made + 1\n            stages.append(nn.Sequential(*blocks))\n        self.stages = nn.Sequential(*stages)\n\n        # --- head ---\n        self.bn_out  = nn.BatchNorm2d(cin)\n        self.dropout = nn.Dropout(dropout)\n        self.fc      = nn.Linear(cin, num_classes)\n        self._init_weights()\n\n    def _init_weights(self):\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                nn.init.kaiming_normal_(m.weight, mode=\"fan_out\", nonlinearity=\"relu\")\n            elif isinstance(m, nn.Conv1d):\n                nn.init.kaiming_normal_(m.weight, nonlinearity=\"relu\")\n            elif isinstance(m, nn.BatchNorm2d):\n                nn.init.ones_(m.weight); nn.init.zeros_(m.bias)\n            elif isinstance(m, nn.Linear):\n                nn.init.trunc_normal_(m.weight, std=0.02); nn.init.zeros_(m.bias)\n\n    def forward(self, x):\n        x = self.stem(x)\n        x = self.stages(x)\n        x = F.silu(self.bn_out(x))\n        x = x.mean(dim=(2, 3))                   # global average pool\n        return self.fc(self.dropout(x))\n\n    def param_groups(self, weight_decay):\n        \"\"\"Split parameters for AdamW. Only the >=2-D tensors get weight decay;\n        BatchNorm scales and shifts, biases and LayerScale are all 1-D and\n        decaying them does measurable harm.\"\"\"\n        decay     = [p for p in self.parameters() if p.requires_grad and p.ndim >= 2]\n        no_decay  = [p for p in self.parameters() if p.requires_grad and p.ndim < 2]\n        return [{\"params\": decay,    \"weight_decay\": weight_decay},\n                {\"params\": no_decay, \"weight_decay\": 0.0}]\n\n\nCONFIGS = {\n    \"small\": dict(widths=(64, 128, 256),  depths=(3, 3, 3), stem_width=32, drop_path=0.05),\n    \"base\":  dict(widths=(128, 256, 512), depths=(4, 4, 4), stem_width=64, drop_path=0.10),\n    \"wide\":  dict(widths=(160, 320, 640), depths=(4, 4, 4), stem_width=64, drop_path=0.15),\n    \"deep\":  dict(widths=(128, 256, 512), depths=(6, 6, 6), stem_width=64, drop_path=0.20),\n}\n\ndef build_model(config=\"base\", **overrides):\n    kw = dict(CONFIGS[config]); kw.update(overrides)\n    return PyraNet(**kw)"},{"cell_type":"code","execution_count":null,"id":"c07","metadata":{},"outputs":[],"source":"for name in CONFIGS:\n    m = build_model(name)\n    n = sum(p.numel() for p in m.parameters())\n    print(f\"{name:6s}  {n/1e6:6.2f}M params   out={tuple(m(torch.randn(2,3,32,32)).shape)}\")\n\nprint()\n_m = build_model(CFG[\"arch\"])\nprint(f\"--- selected: {CFG['arch']} ---\")\n_h = torch.randn(1, 3, 32, 32)\n_h = _m.stem(_h); print(f\"after stem    {tuple(_h.shape)}\")\nfor i, st in enumerate(_m.stages):\n    _h = st(_h);   print(f\"after stage {i+1} {tuple(_h.shape)}\")\nprint(f\"ECA params in whole net: \"\n      f\"{sum(p.numel() for mod in _m.modules() if isinstance(mod, ECA) for p in mod.parameters())}\")\ndel _m, _h"},{"cell_type":"markdown","id":"c08","metadata":{},"source":"---\n## 2 · Data and augmentation\n\n**What's used for what.** Kaggle gives 50 000 labelled training images and a\n300 000-image test set with no labels (10 000 real test images, 290 000 junk\nones that the scoring ignores). There's no labelled test split to borrow here,\nso I carve a class-stratified 5 000 images out of the training set just to get\nan honest accuracy curve, and train on the remaining 45 000. The leaderboard\nis the real held-out measurement.\n\n**Why the whole pipeline sits on the GPU.** CIFAR-10 is about 150 MB as `uint8`,\nwhich fits in VRAM with room to spare. So the dataset gets uploaded once and\nevery augmentation runs as a batched tensor op right there. No `DataLoader`, no\nworker processes, no host-to-device copy in the inner loop — which is the only\nreason a 300-epoch run is a couple of hours rather than most of a day.\n\nEverything below is written from scratch; no `torchvision.transforms`.\n\n| | what | details |\n|---|---|---|\n| geometric | random crop | 32x32 out of a 40x40 reflect-padded canvas |\n| | horizontal flip | p = 0.5 |\n| | random affine | ±12° rotation, ±8° shear, ±10% scale, one batched `grid_sample` |\n| photometric | brightness / contrast / saturation | ±25%, each independently at p = 0.5 |\n| occlusion | **Cutout** | one 12x12 square zeroed out |\n| mixing | **CutMix** | a random box pasted in from a shuffled batch |\n| | MixUp | plain convex combination, the other half of the mixed batches |\n\nCutout zeroes rather than fills with grey because the images are already\nmean-normalised at that point, so zero *is* the dataset mean.\n\nCutMix and MixUp fire on half the batches; of those, half get one and half the\nother. Alternating them works better for me than committing to either: CutMix\npushes the model to find localised evidence, MixUp smooths the decision\nboundary, and they don't seem to get in each other's way."},{"cell_type":"code","execution_count":null,"id":"c09","metadata":{},"outputs":[],"source":"CIFAR_MEAN = (0.4914, 0.4822, 0.4465)\nCIFAR_STD  = (0.2470, 0.2435, 0.2616)\n# Kaggle's trainLabels.csv uses these exact strings, and this order is also the\n# canonical CIFAR-10 class order, so index i below == canonical label i.\nCLASSES = [\"airplane\", \"automobile\", \"bird\", \"cat\", \"deer\",\n           \"dog\", \"frog\", \"horse\", \"ship\", \"truck\"]\n\n\ndef extract_7z(archive, dest):\n    \"\"\"Unpack a .7z with whatever the image provides (no internet on Kaggle).\"\"\"\n    try:\n        import py7zr\n        with py7zr.SevenZipFile(archive, \"r\") as z:\n            z.extractall(path=dest)\n        return \"py7zr\"\n    except ImportError:\n        pass\n    if not os.path.exists(archive):\n        raise FileNotFoundError(f\"archive not found: {archive}\")\n    for exe in (\"7z\", \"7za\", \"7zr\"):\n        if shutil.which(exe) is None:\n            continue\n        r = subprocess.run([exe, \"x\", \"-y\", f\"-o{dest}\", archive],\n                           capture_output=True, text=True)\n        if r.returncode == 0:\n            return exe\n        # Don't swallow the reason the way a bare check=True would.\n        print(f\"{exe} failed (exit {r.returncode}): \"\n              f\"{(r.stderr or r.stdout or '')[-400:]}\")\n    raise RuntimeError(\"could not unpack; tried py7zr, 7z, 7za, 7zr\")\n\n\ndef decode_png_dir(directory, n, cache, workers=16):\n    \"\"\"Decode <directory>/1.png .. n.png into one uint8 NCHW array, cached.\"\"\"\n    if os.path.exists(cache):\n        print(\"using cached\", cache)\n        return np.load(cache, mmap_mode=\"r\")\n    arr = np.zeros((n, 3, 32, 32), dtype=np.uint8)\n\n    def one(i):\n        with Image.open(os.path.join(directory, f\"{i + 1}.png\")) as im:\n            arr[i] = np.asarray(im.convert(\"RGB\"), dtype=np.uint8).transpose(2, 0, 1)\n\n    t0 = time.time()\n    with ThreadPoolExecutor(max_workers=workers) as ex:\n        list(ex.map(one, range(n)))\n    np.save(cache, arr)\n    print(f\"decoded {n} images in {time.time() - t0:.0f}s -> {cache}\")\n    return arr\n\n\ndef load_split(name, n):\n    \"\"\"Extract <name>.7z from the competition mount and decode it.\"\"\"\n    d = os.path.join(SCRATCH, name)\n    if not (os.path.isdir(d) and len(os.listdir(d)) >= n):\n        t0 = time.time()\n        used = extract_7z(os.path.join(COMP_DIR, f\"{name}.7z\"), SCRATCH)\n        print(f\"extracted {name}.7z with {used} in {time.time() - t0:.0f}s\")\n    got = len(os.listdir(d))\n    assert got == n, f\"expected {n} files in {d}, found {got}\"\n    return decode_png_dir(d, n, os.path.join(SCRATCH, f\"{name}_u8.npy\"))\n\n\n# ---- labelled training data (50 000) ----\nx_all_np = load_split(\"train\", 50000)\n_lab = pd.read_csv(os.path.join(COMP_DIR, \"trainLabels.csv\")).sort_values(\"id\")\nassert _lab[\"id\"].tolist() == list(range(1, 50001)), \"trainLabels.csv ids must be 1..50000\"\nassert set(_lab[\"label\"]) == set(CLASSES), set(_lab[\"label\"]) ^ set(CLASSES)\ny_all = torch.tensor([CLASSES.index(s) for s in _lab[\"label\"]], dtype=torch.long)\nx_all = torch.from_numpy(np.ascontiguousarray(x_all_np))\nprint(\"labelled set:\", tuple(x_all.shape), \"| class counts:\", y_all.bincount().tolist())\n\n# ---- train / validation split -------------------------------------------\n# Unlike the offline version there is no labelled test set here: the only\n# labels Kaggle gives are these 50 000. So a fixed, class-stratified slice is\n# held out purely to produce an honest accuracy curve; the leaderboard over the\n# 300 000-image test set is the real held-out metric.\nVAL_N = CFG[\"val_n\"]\n_g = torch.Generator().manual_seed(CFG[\"seed\"])\n_val_parts = []\nfor c in range(10):                              # stratified: VAL_N/10 per class\n    idx = (y_all == c).nonzero(as_tuple=True)[0]\n    _val_parts.append(idx[torch.randperm(len(idx), generator=_g)[: VAL_N // 10]])\nval_idx = torch.cat(_val_parts)\nmask = torch.ones(len(y_all), dtype=torch.bool); mask[val_idx] = False\ntr_idx = mask.nonzero(as_tuple=True)[0]\n\nx_train, y_train = x_all[tr_idx], y_all[tr_idx]\nx_val,   y_val   = x_all[val_idx], y_all[val_idx]\nprint(f\"train {tuple(x_train.shape)} counts {y_train.bincount().tolist()}\")\nprint(f\"val   {tuple(x_val.shape)} counts {y_val.bincount().tolist()}\")"},{"cell_type":"code","execution_count":null,"id":"c11","metadata":{},"outputs":[],"source":"class GPUBatches:\n    \"\"\"Keeps a uint8 dataset in VRAM and hands back augmented float batches.\"\"\"\n\n    def __init__(self, x_u8, y, device, batch_size, pad=4, train=True, drop_last=None):\n        self.device, self.train, self.bs, self.pad = device, train, batch_size, pad\n        self.drop_last = train if drop_last is None else drop_last\n        self.y = y.to(device)\n        self.n = x_u8.shape[0]\n\n        x = x_u8.to(device)\n        # Pad once up front so the per-batch random crop is a pure gather.\n        self.x = F.pad(x.float(), (pad,) * 4, mode=\"reflect\").to(torch.uint8) if train else x\n\n        self.mean = torch.tensor(CIFAR_MEAN, device=device).view(1, 3, 1, 1) * 255.0\n        self.std  = torch.tensor(CIFAR_STD,  device=device).view(1, 3, 1, 1) * 255.0\n        self._ar  = torch.arange(32, device=device)\n\n    # ---------------- individual augmentations ---------------- #\n    def _crop_flip(self, imgs):\n        b, span = imgs.shape[0], 2 * self.pad + 1\n        oy = torch.randint(0, span, (b,), device=self.device)\n        ox = torch.randint(0, span, (b,), device=self.device)\n        rows = (oy[:, None] + self._ar[None, :])[:, None, :, None]   # b,1,32,1\n        cols = (ox[:, None] + self._ar[None, :])[:, None, None, :]   # b,1,1,32\n        bi = torch.arange(b, device=self.device)[:, None, None, None]\n        ci = torch.arange(3, device=self.device)[None, :, None, None]\n        out = imgs[bi, ci, rows, cols]                               # b,3,32,32\n        flip = torch.rand(b, device=self.device) < 0.5\n        return torch.where(flip[:, None, None, None], out.flip(-1), out)\n\n    def _affine(self, imgs, p=0.5, deg=12.0, shear=8.0, scale=0.1):\n        \"\"\"Rotate, shear and scale, all folded into one batched grid_sample.\"\"\"\n        b = imgs.shape[0]\n        use = (torch.rand(b, device=self.device) < p).float()\n        rot = torch.empty(b, device=self.device).uniform_(-deg, deg)   * (math.pi / 180) * use\n        shx = torch.empty(b, device=self.device).uniform_(-shear, shear) * (math.pi / 180) * use\n        sc  = torch.empty(b, device=self.device).uniform_(1 - scale, 1 + scale) * use + (1 - use)\n\n        cos, sin, tan = torch.cos(rot), torch.sin(rot), torch.tan(shx)\n        theta = torch.zeros(b, 2, 3, device=self.device)   # rotation @ shear, / scale\n        theta[:, 0, 0] = cos / sc\n        theta[:, 0, 1] = (cos * tan - sin) / sc\n        theta[:, 1, 0] = sin / sc\n        theta[:, 1, 1] = (sin * tan + cos) / sc\n        grid = F.affine_grid(theta, list(imgs.shape), align_corners=False)\n        return F.grid_sample(imgs, grid, mode=\"bilinear\",\n                             padding_mode=\"reflection\", align_corners=False)\n\n    def _color(self, imgs, p=0.5, strength=0.25):\n        b = imgs.shape[0]\n        def f():\n            v = torch.empty(b, 1, 1, 1, device=self.device).uniform_(1 - strength, 1 + strength)\n            keep = (torch.rand(b, 1, 1, 1, device=self.device) < p).float()\n            return v * keep + (1 - keep)\n        imgs = imgs * f()                                             # brightness\n        mu = imgs.mean(dim=(1, 2, 3), keepdim=True)\n        imgs = (imgs - mu) * f() + mu                                 # contrast\n        gray = imgs.mean(dim=1, keepdim=True)\n        return (imgs - gray) * f() + gray                             # saturation\n\n    def _cutout(self, imgs, size=12):\n        \"\"\"Cutout. Zeroing rather than greying out is correct here, not lazy:\n        the batch is already mean-normalised, so zero is the dataset mean.\"\"\"\n        b, _, h, w = imgs.shape\n        cy = torch.randint(0, h, (b,), device=self.device)\n        cx = torch.randint(0, w, (b,), device=self.device)\n        r = self._ar[:h]\n        my = (r[None, :] >= cy[:, None] - size // 2) & (r[None, :] < cy[:, None] + size // 2)\n        mx = (r[None, :] >= cx[:, None] - size // 2) & (r[None, :] < cx[:, None] + size // 2)\n        return imgs.masked_fill((my[:, :, None] & mx[:, None, :]).unsqueeze(1), 0.0)\n\n    def _normalize(self, u8):\n        return (u8.float() - self.mean) / self.std\n\n    # ---------------- iteration ---------------- #\n    def epoch(self, aug=None):\n        perm = (torch.randperm(self.n, device=self.device) if self.train\n                else torch.arange(self.n, device=self.device))\n        stop = (self.n // self.bs) * self.bs if self.drop_last else self.n\n        for i in range(0, stop, self.bs):\n            idx = perm[i:i + self.bs]\n            imgs, y = self.x[idx], self.y[idx]\n            if self.train:\n                imgs = self._normalize(self._crop_flip(imgs))\n                if aug.get(\"affine_p\", 0) > 0:\n                    imgs = self._affine(imgs, p=aug[\"affine_p\"])\n                if aug.get(\"color_p\", 0) > 0:\n                    imgs = self._color(imgs, p=aug[\"color_p\"], strength=aug[\"color_strength\"])\n                if aug.get(\"cutout\", 0) > 0:\n                    imgs = self._cutout(imgs, size=aug[\"cutout\"])\n            else:\n                imgs = self._normalize(imgs)\n            yield imgs.contiguous(memory_format=torch.channels_last), y\n\n    def __len__(self):\n        return self.n // self.bs if self.drop_last else (self.n + self.bs - 1) // self.bs"},{"cell_type":"code","execution_count":null,"id":"c12","metadata":{},"outputs":[],"source":"def rand_bbox(h, w, lam, device):\n    \"\"\"A box whose area fraction is (1 - lam), clipped to the image.\"\"\"\n    cut = math.sqrt(1.0 - lam)\n    ch, cw = int(h * cut), int(w * cut)\n    cy = int(torch.randint(0, h, (1,), device=device))\n    cx = int(torch.randint(0, w, (1,), device=device))\n    return (max(cy - ch // 2, 0), min(cy + ch // 2, h),\n            max(cx - cw // 2, 0), min(cx + cw // 2, w))\n\n\ndef mix_batch(x, y, cutmix_alpha=1.0, mixup_alpha=0.2, prob=0.5, cutmix_share=0.5):\n    \"\"\"Apply CutMix or MixUp to a whole batch.\n\n    Returns (x, y_a, y_b, lam). When nothing was mixed, y_a and y_b are the same\n    and lam is 1.0, so the blended loss upstream quietly becomes ordinary\n    cross-entropy and the caller needs no special case.\n    \"\"\"\n    if prob <= 0 or torch.rand(()).item() > prob:\n        return x, y, y, 1.0\n\n    perm = torch.randperm(x.shape[0], device=x.device)\n\n    if torch.rand(()).item() < cutmix_share:\n        # ---- CutMix: paste a rectangle from the shuffled batch ----\n        lam = float(np.random.beta(cutmix_alpha, cutmix_alpha))\n        h, w = x.shape[2], x.shape[3]\n        y1, y2, x1, x2 = rand_bbox(h, w, lam, x.device)\n        x = x.clone()\n        x[:, :, y1:y2, x1:x2] = x[perm][:, :, y1:y2, x1:x2]\n        lam = 1.0 - (y2 - y1) * (x2 - x1) / (h * w)   # exact pasted area\n    else:\n        # ---- MixUp: convex combination ----\n        lam = float(np.random.beta(mixup_alpha, mixup_alpha))\n        x = lam * x + (1.0 - lam) * x[perm]\n\n    return x, y, y[perm], lam"},{"cell_type":"code","execution_count":null,"id":"c13","metadata":{},"outputs":[],"source":"# --- sanity check: visualise one augmented batch -------------------------\nimport matplotlib.pyplot as plt\n\n_vis = GPUBatches(x_train[:16], y_train[:16], device, 16, pad=CFG[\"pad\"], train=True, drop_last=False)\n_xb, _yb = next(_vis.epoch(CFG))\n_xm, _ya, _yb2, _lam = mix_batch(_xb, _yb, CFG[\"cutmix_alpha\"], CFG[\"mixup_alpha\"],\n                                 prob=1.0, cutmix_share=1.0)   # force CutMix for the figure\n\n_mean = torch.tensor(CIFAR_MEAN, device=device).view(1, 3, 1, 1)\n_std  = torch.tensor(CIFAR_STD,  device=device).view(1, 3, 1, 1)\ndef _show(ax, t):\n    ax.imshow((t * _std[0] + _mean[0]).clamp(0, 1).permute(1, 2, 0).cpu().numpy())\n    ax.axis(\"off\")\n\nfig, axes = plt.subplots(2, 8, figsize=(13, 3.6))\nfor i in range(8):\n    _show(axes[0, i], _xb[i]);  axes[0, i].set_title(CLASSES[_yb[i]], fontsize=8)\n    _show(axes[1, i], _xm[i]);  axes[1, i].set_title(f\"{CLASSES[_ya[i]]}/{CLASSES[_yb2[i]]}\", fontsize=8)\naxes[0, 0].set_ylabel(\"geo+colour+cutout\"); axes[1, 0].set_ylabel(\"+CutMix\")\nfig.suptitle(f\"top: crop/flip/affine/colour/Cutout   |   bottom: + CutMix (lam={_lam:.2f})\")\nplt.tight_layout(); plt.show()\ndel _vis, _xb, _yb, _xm"},{"cell_type":"markdown","id":"c14","metadata":{},"source":"---\n## 3 · Optimizer, schedule, and a weight average\n\n**AdamW.** Decoupled weight decay of 0.05, applied only to tensors with two or\nmore dimensions — conv kernels and the classifier matrix. The 1-D parameters\n(BatchNorm scales and shifts, biases, the LayerScale vectors) are left alone,\nbecause decaying a normalisation scale just quietly shrinks the features and\ncosts accuracy for no benefit. I also run `β₂ = 0.99` rather than the usual\n0.999; the gradient scale moves quickly in a from-scratch CNN and the shorter\nsecond-moment horizon tracks it better.\n\n**Linear warmup into cosine annealing,** stepped every iteration rather than\nevery epoch so the curve is actually smooth:\n\n$$\n\\eta(t)=\\begin{cases}\n\\eta_{\\text{peak}}\\dfrac{t+1}{T_{w}} & t < T_{w}\\\\[2ex]\n\\eta_{\\min}+\\tfrac12\\left(\\eta_{\\text{peak}}-\\eta_{\\min}\\right)\n\\left(1+\\cos\\pi\\dfrac{t-T_{w}}{T-T_{w}}\\right) & t \\ge T_{w}\n\\end{cases}\n$$\n\nThe warmup isn't decorative. LayerScale means the residual branches start near\nzero, Adam's variance estimates are still junk for the first few hundred steps,\nand going straight in at 2e-3 blows the stem apart. Five epochs of ramp fixes it.\n\n**An exponential moving average of the weights,** decay 0.9998, with the decay\neased in so the average isn't dragged around by the random initialisation. Under\na long cosine schedule this is close to free accuracy — usually a few tenths of\na percent over the final weights — so I keep it alongside the live model and\nreport both."},{"cell_type":"code","execution_count":null,"id":"c15","metadata":{},"outputs":[],"source":"def lr_at(step, total_steps, warmup_steps, peak_lr, min_lr):\n    \"\"\"Linear warmup then cosine annealing, evaluated per optimiser step.\"\"\"\n    if step < warmup_steps:\n        return peak_lr * (step + 1) / warmup_steps\n    p = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n    return min_lr + 0.5 * (peak_lr - min_lr) * (1.0 + math.cos(math.pi * p))\n\n\n# --- plot the exact schedule this run will use ---\n_spe   = len(GPUBatches(x_train, y_train, torch.device(\"cpu\"), CFG[\"batch_size\"], train=True))\n_total = _spe * CFG[\"epochs\"]\n_warm  = _spe * CFG[\"warmup_epochs\"]\n_minlr = CFG[\"peak_lr\"] * CFG[\"min_lr_frac\"]\n_ys = [lr_at(s, _total, _warm, CFG[\"peak_lr\"], _minlr) for s in range(_total)]\n\nfig, ax = plt.subplots(figsize=(7, 2.8))\nax.plot(np.arange(_total) / _spe, _ys)\nax.axvspan(0, CFG[\"warmup_epochs\"], color=\"orange\", alpha=.25)\nax.annotate(\"linear warmup\", (CFG[\"warmup_epochs\"], CFG[\"peak_lr\"]),\n            xytext=(CFG[\"warmup_epochs\"] + CFG[\"epochs\"] * .04, CFG[\"peak_lr\"] * .95), fontsize=9)\nax.set_xlabel(\"epoch\"); ax.set_ylabel(\"learning rate\")\nax.set_title(f\"AdamW LR: {CFG['warmup_epochs']}-epoch warmup -> cosine to {_minlr:.2e}\"\n             f\"   ({_spe} steps/epoch, {_total} total)\")\nplt.tight_layout(); plt.show()"},{"cell_type":"code","execution_count":null,"id":"c16","metadata":{},"outputs":[],"source":"class EMA:\n    \"\"\"Exponential moving average of the model weights.\n\n    The update runs on every optimiser step, so it has to be cheap. A naive\n    version loops over ~200 state_dict tensors per step, and the kernel-launch\n    overhead alone then dominates the step time. Here the live tensors are\n    cached once and the whole average is applied with two fused `_foreach_`\n    calls -- a couple of kernels regardless of model size.\n    \"\"\"\n\n    def __init__(self, model, decay=0.9998):\n        self.decay, self.n = decay, 0\n        # Cache references to the live tensors. AdamW updates parameters\n        # in-place and the forward pass updates BN buffers in-place, so these\n        # references stay valid for the whole run.\n        self.f_keys, self.f_live, self.f_shadow = [], [], []\n        self.o_keys, self.o_live = [], []\n        for k, v in model.state_dict().items():\n            if v.dtype.is_floating_point:\n                self.f_keys.append(k); self.f_live.append(v)\n                self.f_shadow.append(v.detach().clone().float())\n            else:\n                self.o_keys.append(k); self.o_live.append(v)   # num_batches_tracked\n\n    @torch.no_grad()\n    def update(self):\n        self.n += 1\n        # Ramp the decay in, so early steps are not dominated by the random init.\n        d = min(self.decay, (1 + self.n) / (10 + self.n))\n        torch._foreach_mul_(self.f_shadow, d)\n        torch._foreach_add_(self.f_shadow, self.f_live, alpha=1.0 - d)\n\n    def state_dict(self):\n        sd = {k: v.to(l.dtype) for k, v, l in zip(self.f_keys, self.f_shadow, self.f_live)}\n        sd.update({k: v.clone() for k, v in zip(self.o_keys, self.o_live)})\n        return sd"},{"cell_type":"markdown","id":"c17","metadata":{},"source":"---\n## 4 · Training\n\nMixed precision with the dtype chosen to match the card, `channels_last`, and\n`torch.compile`. Batches are a fixed size so the compiled graph never has to\nre-specialise on a short final batch.\n\nThe loss is the standard CutMix/MixUp form: label-smoothed cross-entropy against\nboth sets of labels, weighted by how much of the image came from each.\n\n$$\\mathcal{L} = \\lambda\\,\\mathrm{CE}(\\hat y, y_a) + (1-\\lambda)\\,\\mathrm{CE}(\\hat y, y_b)$$"},{"cell_type":"code","execution_count":null,"id":"c18","metadata":{},"outputs":[],"source":"@torch.no_grad()\ndef evaluate(model, batches, tta=False):\n    \"\"\"Top-1 accuracy. tta=True also averages in the horizontal flip.\"\"\"\n    model.eval()\n    correct = total = 0\n    for x, y in batches.epoch():\n        with torch.autocast(\"cuda\", dtype=AMP_DTYPE):\n            logits = model(x).float()\n            if tta:\n                logits = logits + model(x.flip(-1)).float()\n        correct += (logits.argmax(1) == y).sum().item()\n        total   += y.numel()\n    return correct / total\n\n\n@torch.no_grad()\ndef predict_logits(model, x_u8, batch_size=1024, tta=True):\n    \"\"\"Class probabilities for a uint8 NCHW array, streamed to the GPU in chunks\n    so the 300k test set doesn't have to fit in VRAM at once.\"\"\"\n    model.eval()\n    mean = torch.tensor(CIFAR_MEAN, device=device).view(1, 3, 1, 1) * 255.0\n    std  = torch.tensor(CIFAR_STD,  device=device).view(1, 3, 1, 1) * 255.0\n    out = []\n    for i in range(0, len(x_u8), batch_size):\n        chunk = x_u8[i:i + batch_size]\n        if not torch.is_tensor(chunk):\n            chunk = torch.from_numpy(np.ascontiguousarray(chunk))\n        xb = ((chunk.to(device, non_blocking=True).float() - mean) / std)\n        xb = xb.contiguous(memory_format=torch.channels_last)\n        with torch.autocast(\"cuda\", dtype=AMP_DTYPE):\n            p = model(xb).float().softmax(1)\n            if tta:\n                p = p + model(xb.flip(-1)).float().softmax(1)\n        out.append((p / (2 if tta else 1)).cpu())\n    return torch.cat(out)"},{"cell_type":"code","execution_count":null,"id":"c19","metadata":{},"outputs":[],"source":"def train(cfg, x_tr, y_tr, x_ev, y_ev, tag=\"run\", log_every=10):\n    torch.manual_seed(cfg[\"seed\"]); np.random.seed(cfg[\"seed\"])\n\n    train_b = GPUBatches(x_tr, y_tr, device, cfg[\"batch_size\"], pad=cfg[\"pad\"], train=True)\n    eval_b  = GPUBatches(x_ev, y_ev, device, 1024, train=False)\n\n    model = build_model(cfg[\"arch\"]).to(device, memory_format=torch.channels_last)\n    n_par = sum(p.numel() for p in model.parameters())\n\n    opt = torch.optim.AdamW(model.param_groups(cfg[\"weight_decay\"]),\n                            lr=cfg[\"peak_lr\"], betas=cfg[\"betas\"], eps=1e-8)\n    crit = nn.CrossEntropyLoss(label_smoothing=cfg[\"label_smooth\"])\n    ema  = EMA(model, cfg[\"ema_decay\"])\n\n    scaler = torch.amp.GradScaler(\"cuda\", enabled=(AMP_DTYPE is torch.float16))\n\n    net = model\n    if cfg[\"compile_model\"]:\n        # Compile, then actually run one step through the compiled graph. A\n        # backend that fails only at runtime would otherwise take the whole run\n        # down hours later; this way it falls back to eager immediately.\n        try:\n            cand = torch.compile(model)\n            probe = torch.zeros(cfg[\"batch_size\"], 3, 32, 32, device=device\n                                ).contiguous(memory_format=torch.channels_last)\n            with torch.autocast(\"cuda\", dtype=AMP_DTYPE):\n                cand(probe).float().sum().backward()\n            model.zero_grad(set_to_none=True)\n            del probe\n            net = cand\n            print(\"torch.compile enabled\")\n        except Exception as e:\n            print(f\"torch.compile unavailable ({type(e).__name__}: {e}); running eager\")\n            net = model\n\n    spe    = len(train_b)\n    total  = spe * cfg[\"epochs\"]\n    warm   = spe * cfg[\"warmup_epochs\"]\n    min_lr = cfg[\"peak_lr\"] * cfg[\"min_lr_frac\"]\n\n    print(f\"{tag}: {cfg['arch']} ({n_par/1e6:.2f}M params), \"\n          f\"{cfg['epochs']} epochs x {spe} steps = {total} steps\")\n\n    # Per-epoch progress is appended to a jsonl file as well as printed, so a\n    # long run can be followed from outside the notebook (tail -f).\n    log_path = os.path.join(OUT_DIR, f\"{tag}_log.jsonl\")\n    open(log_path, \"w\").close()\n\n    hist, step, t0 = [], 0, time.time()\n    for ep in range(cfg[\"epochs\"]):\n        model.train()\n        run_loss, seen = torch.zeros((), device=device), 0\n        for x, y in train_b.epoch(cfg):\n            lr = lr_at(step, total, warm, cfg[\"peak_lr\"], min_lr)\n            for g in opt.param_groups:\n                g[\"lr\"] = lr\n\n            x, ya, yb, lam = mix_batch(x, y, cfg[\"cutmix_alpha\"], cfg[\"mixup_alpha\"],\n                                       cfg[\"mix_prob\"], cfg[\"cutmix_share\"])\n            with torch.autocast(\"cuda\", dtype=AMP_DTYPE):\n                logits = net(x)\n                loss = lam * crit(logits, ya) + (1.0 - lam) * crit(logits, yb)\n\n            opt.zero_grad(set_to_none=True)\n            scaler.scale(loss).backward()      # no-op scaling under bf16\n            scaler.step(opt)\n            scaler.update()\n            ema.update()\n\n            # Accumulate on-device: .item() here would force a host sync\n            # every single step and stall the pipeline.\n            run_loss += loss.detach() * y.numel(); seen += y.numel(); step += 1\n\n        rec = {\"epoch\": ep + 1, \"loss\": run_loss.item() / seen, \"lr\": lr,\n               \"time\": time.time() - t0}\n        if (ep + 1) % log_every == 0 or ep + 1 == cfg[\"epochs\"] or ep < 3:\n            rec[\"val_acc\"] = evaluate(net, eval_b)\n            ema_model = build_model(cfg[\"arch\"]).to(device, memory_format=torch.channels_last)\n            ema_model.load_state_dict(ema.state_dict())\n            rec[\"ema_acc\"] = evaluate(ema_model, eval_b)\n            del ema_model\n            print(f\"ep {ep+1:3d}/{cfg['epochs']}  loss {rec['loss']:.4f}  lr {lr:.2e}  \"\n                  f\"acc {rec['val_acc']*100:.2f}%  ema {rec['ema_acc']*100:.2f}%  \"\n                  f\"[{rec['time']/60:.1f} min]\")\n        hist.append(rec)\n        with open(log_path, \"a\") as fh:\n            fh.write(json.dumps(rec) + \"\\n\")\n\n    # final EMA model\n    final_ema = build_model(cfg[\"arch\"]).to(device, memory_format=torch.channels_last)\n    final_ema.load_state_dict(ema.state_dict())\n\n    ckpt = os.path.join(OUT_DIR, f\"{tag}.pt\")\n    torch.save({\"cfg\": cfg, \"model\": model.state_dict(),\n                \"ema\": ema.state_dict(), \"hist\": hist}, ckpt)\n    print(f\"saved {ckpt}  ({time.time()-t0:.0f}s total)\")\n    return model, final_ema, hist"},{"cell_type":"code","execution_count":null,"id":"c20","metadata":{},"outputs":[],"source":"# Training takes ~2.8 h for the default 300-epoch config on an RTX 3080, so the\n# cell is idempotent: if the checkpoint is already there it is loaded instead,\n# which makes re-running the notebook (e.g. to regenerate the submission)\n# cheap. Force a fresh run with RETRAIN=1.\nckpt_path = os.path.join(OUT_DIR, f\"pyranet_{CFG['arch']}.pt\")\n\nif os.path.exists(ckpt_path) and os.environ.get(\"RETRAIN\") != \"1\":\n    print(f\"loading {ckpt_path}  (set RETRAIN=1 to train from scratch)\")\n    ck = torch.load(ckpt_path, map_location=device, weights_only=False)\n    model, ema_model = (build_model(ck[\"cfg\"][\"arch\"]).to(device, memory_format=torch.channels_last)\n                        for _ in range(2))\n    model.load_state_dict(ck[\"model\"])\n    ema_model.load_state_dict(ck[\"ema\"])\n    hist = ck[\"hist\"]\n    print(f\"restored after {hist[-1]['epoch']} epochs, {hist[-1]['time']/60:.0f} min of training\")\nelse:\n    model, ema_model, hist = train(CFG, x_train, y_train, x_val, y_val,\n                                   tag=f\"pyranet_{CFG['arch']}\")"},{"cell_type":"markdown","id":"c21","metadata":{},"source":"---\n## 5 · How it did"},{"cell_type":"code","execution_count":null,"id":"c22","metadata":{},"outputs":[],"source":"acc_raw      = evaluate(model,     GPUBatches(x_val, y_val, device, 1024, train=False))\nacc_raw_tta  = evaluate(model,     GPUBatches(x_val, y_val, device, 1024, train=False), tta=True)\nacc_ema      = evaluate(ema_model, GPUBatches(x_val, y_val, device, 1024, train=False))\nacc_ema_tta  = evaluate(ema_model, GPUBatches(x_val, y_val, device, 1024, train=False), tta=True)\n\nprint(f\"{'final weights':<22} {acc_raw*100:6.2f}%\")\nprint(f\"{'final weights + flip':<22} {acc_raw_tta*100:6.2f}%\")\nprint(f\"{'EMA weights':<22} {acc_ema*100:6.2f}%\")\nprint(f\"{'EMA weights + flip':<22} {acc_ema_tta*100:6.2f}%   <- used for the submission\")\n\nGRADE = [(0.942, 6), (0.931, 5), (0.921, 4), (0.911, 3), (0.901, 2), (0.88, 1)]\npts = next((p for t, p in GRADE if acc_ema_tta >= t), 0)\nprint(f\"\\nheld-out validation {acc_ema_tta:.4f}; the leaderboard score over the\\n\"\n      f\"300k test set is the figure that counts  ->  indicatively {pts}/6\")"},{"cell_type":"code","execution_count":null,"id":"c23","metadata":{},"outputs":[],"source":"ep   = [h[\"epoch\"] for h in hist]\nloss = [h[\"loss\"]  for h in hist]\nev   = [(h[\"epoch\"], h[\"val_acc\"], h[\"ema_acc\"]) for h in hist if \"val_acc\" in h]\n\nfig, axes = plt.subplots(1, 3, figsize=(15, 3.6))\naxes[0].plot(ep, loss);                      axes[0].set_title(\"training loss (CutMix/MixUp)\")\naxes[1].plot([e for e,_,_ in ev], [a*100 for _,a,_ in ev], label=\"raw\")\naxes[1].plot([e for e,_,_ in ev], [a*100 for _,_,a in ev], label=\"EMA\")\naxes[1].axhline(94.2, ls=\"--\", c=\"r\", lw=1, label=\"94.2% (6/6)\")\naxes[1].legend(); axes[1].set_title(\"validation accuracy (%)\"); axes[1].set_ylim(60, 100)\naxes[2].plot(ep, [h[\"lr\"] for h in hist]);   axes[2].set_title(\"learning rate\")\nfor a in axes: a.set_xlabel(\"epoch\"); a.grid(alpha=.3)\nplt.tight_layout(); plt.show()"},{"cell_type":"code","execution_count":null,"id":"c24","metadata":{},"outputs":[],"source":"# per-class accuracy and the worst confusions\nprobs = predict_logits(ema_model, x_val, tta=True)\npred  = probs.argmax(1)\ncm = torch.zeros(10, 10, dtype=torch.long)\nfor t, p in zip(y_val, pred):\n    cm[t, p] += 1\n\nprint(\"per-class accuracy\")\nfor i, c in enumerate(CLASSES):\n    print(f\"  {c:<11} {cm[i,i].item()/cm[i].sum().item()*100:6.2f}%\")\n\noff = [(cm[i, j].item(), CLASSES[i], CLASSES[j])\n       for i in range(10) for j in range(10) if i != j]\nprint(\"\\ntop confusions (true -> predicted)\")\nfor n, a, b in sorted(off, reverse=True)[:6]:\n    print(f\"  {a:<11} -> {b:<11} {n}\")"},{"cell_type":"markdown","id":"c25","metadata":{},"source":"---\n## 6 · Kaggle submission\n\nThe competition archives are already mounted read-only at\n`/kaggle/input/cifar-10`, so the test set is unpacked straight from there.\nPredictions for all 300 000 images are written to\n`/kaggle/working/submission.csv`, which can be submitted from this\nnotebook's output tab."},{"cell_type":"code","execution_count":null,"id":"c26","metadata":{},"outputs":[],"source":"# The competition test set: 300 000 PNGs, of which 10 000 are the real\n# CIFAR-10 test images and the rest are junk that Kaggle ignores when scoring.\ntry:\n    x_kaggle = load_split(\"test\", 300000)\n    print(\"kaggle test:\", x_kaggle.shape)\nexcept Exception as e:\n    x_kaggle = None\n    print(\"SKIPPING submission:\", type(e).__name__, e)"},{"cell_type":"code","execution_count":null,"id":"c27","metadata":{},"outputs":[],"source":"if x_kaggle is not None:\n    t0 = time.time()\n    probs = predict_logits(ema_model, x_kaggle, batch_size=2048, tta=True)\n    labels = [CLASSES[i] for i in probs.argmax(1).tolist()]\n    print(f\"predicted {len(labels)} images in {time.time() - t0:.0f}s\")\n\n    sub_path = os.path.join(OUT_DIR, \"submission.csv\")\n    sub = pd.DataFrame({\"id\": np.arange(1, len(labels) + 1), \"label\": labels})\n    sub.to_csv(sub_path, index=False)\n\n    # Shape checks that would catch a mis-ordered or truncated submission.\n    assert len(sub) == 300000 and sub[\"id\"].iloc[0] == 1 and sub[\"id\"].iloc[-1] == 300000\n    assert set(sub[\"label\"]) <= set(CLASSES)\n    print(sub.head())\n    print(\"\\npredicted label distribution over all 300k (incl. junk):\")\n    print(sub[\"label\"].value_counts().sort_index().to_string())\n    print(\"\\nwrote\", sub_path)"},{"cell_type":"markdown","id":"c28","metadata":{},"source":"---\n## 7 · Summary\n\n**The network.** PyraNet-ECA: a three-stage residual pyramid, deep stem at full\nresolution, pre-activation blocks carrying ECA channel attention, residual\nbranches gated by LayerScale, stochastic depth ramped with depth, and\nanti-aliased `AvgPool -> 1x1` shortcuts on the downsampling paths. Assembled\nfrom `nn.Conv2d` and `nn.BatchNorm2d` in §1. No pre-trained weights anywhere,\nand no model library.\n\n**Optimizer.** AdamW, decoupled decay 0.05 on the ≥2-D tensors only,\n`β = (0.9, 0.99)`.\n\n**Schedule.** Five epochs of linear warmup to 2e-3, then cosine down to 2e-6,\nstepped per iteration.\n\n**Augmentation.** CutMix alternating with MixUp on half the batches, over a base\nof random crop, flip, affine, colour jitter and Cutout — all hand-written as\nbatched GPU ops.\n\n**At test time.** EMA weights, with the horizontal flip averaged in."}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.3"}},"nbformat":4,"nbformat_minor":5}