{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10"},"kaggle":{"accelerator":"nvidiaTeslaT4","isInternetEnabled":true,"isGpuEnabled":true,"language":"python","sourceType":"notebook"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CIFAR-10 Object Recognition: Custom SE-PreAct Residual Network Trained from Scratch\n\n**Name:** Ali Khalid\n**Course:** Applied Deep Learning\n**Competition:** [CIFAR-10 - Object Recognition in Images](https://www.kaggle.com/c/cifar-10)\n\n\nNo pre-trained weights are used and nothing is imported from `torchvision.models`; all weights start from random (Kaiming) initialisation.\n\n\n\n","metadata":{}},{"cell_type":"markdown","source":"## Step 1: Environment setup and fast extraction of the test set\n\n`test.7z` holds 300,000 PNG files. Writing that many small files to the Kaggle disk is very slow, so the archive is read as a stream\nwith `libarchive` and every PNG is decoded straight into one NumPy array of shape (300000, 32, 32, 3). Row *k* of the array is image id *k*+1.\nIf `libarchive` cannot be loaded, the code falls back to normal extraction (7z, bsdtar or py7zr) into the RAM disk and counts the extracted files.","metadata":{}},{"cell_type":"code","source":"# ---- Step 1.1: imports, random seeds and GPU check ----\nimport os, io, sys, glob, time, math, random, shutil, subprocess, importlib, warnings\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nfrom torch.utils.data import Dataset, DataLoader\n\n# The mixed-precision scaler may skip the first optimizer step, which triggers a harmless scheduler warning.\nwarnings.filterwarnings(\"ignore\", message=\".*lr_scheduler.step.*\")\n\nSEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\ntorch.cuda.manual_seed_all(SEED)\ntorch.backends.cudnn.benchmark = True      # input size never changes, so let cuDNN pick the fastest kernels\n\nINPUT_DIR = \"/kaggle/input\"\nWORK_DIR = \"/kaggle/working\"\nos.makedirs(WORK_DIR, exist_ok=True)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"PyTorch\", torch.__version__, \"| torchvision\", torchvision.__version__, flush=True)\n\nif device.type != \"cuda\":\n    raise RuntimeError(\"No GPU found. Set Settings > Accelerator to 'GPU T4 x2' and run again.\")\n\nn_gpus = torch.cuda.device_count()\nfor i in range(n_gpus):\n    print(f\"GPU {i}: {torch.cuda.get_device_name(i)}\", flush=True)\n\n# Recent PyTorch builds no longer include kernels for old GPUs such as the P100 (sm_60).\nmajor, minor = torch.cuda.get_device_capability(0)\nif major < 7 and f\"sm_{major}{minor}\" not in torch.cuda.get_arch_list():\n    raise RuntimeError(\"This PyTorch build does not support the selected GPU. \"\n                       \"Set Settings > Accelerator to 'GPU T4 x2' and run again.\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- Step 1.2: locate the competition files ----\ndef find_file(name):\n    # The folder name under /kaggle/input can differ between accounts, so search for the file.\n    hits = sorted(glob.glob(os.path.join(INPUT_DIR, \"**\", name), recursive=True))\n    if not hits:\n        raise FileNotFoundError(f\"{name} was not found. Add the CIFAR-10 competition data to this notebook.\")\n    return hits[0]\n\nTEST_7Z = find_file(\"test.7z\")\nSAMPLE_SUB = find_file(\"sampleSubmission.csv\")\nN_TEST = len(pd.read_csv(SAMPLE_SUB))          # 300,000 rows, one per test image\n\nprint(f\"test.7z: {TEST_7Z} ({os.path.getsize(TEST_7Z) / 1e6:.0f} MB)\", flush=True)\nprint(f\"sampleSubmission.csv: {N_TEST:,} rows\", flush=True)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- Step 1.3: functions that read a .7z archive of PNG images ----\ndef pip_install(package):\n    subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", package], check=False)\n    importlib.invalidate_caches()               # lets Python see a package installed while the notebook runs\n\ndef load_libarchive():\n    # Plan A needs the libarchive-c package and the system libarchive library.\n    for attempt in range(3):\n        try:\n            import libarchive\n            return libarchive\n        except Exception as error:\n            last_error = error\n        if attempt == 0:\n            pip_install(\"libarchive-c\")\n        elif attempt == 1:\n            # \"apt-get update\" must come first, otherwise the package download can fail with 404.\n            subprocess.run(\"apt-get update -qq && apt-get install -y -qq libarchive-dev > /dev/null\",\n                           shell=True, check=False)\n    print(\"libarchive could not be loaded:\", last_error, flush=True)\n    return None\n\ndef extract_to_folder(archive_path):\n    # Plan B: normal extraction. /dev/shm is a RAM disk, which handles many small files much better than the normal disk.\n    base = \"/dev/shm\" if os.path.isdir(\"/dev/shm\") and shutil.disk_usage(\"/dev/shm\").free > 3e9 else \"/tmp\"\n    out_dir = os.path.join(base, \"extract_\" + os.path.basename(archive_path).split(\".\")[0])\n    shutil.rmtree(out_dir, ignore_errors=True)\n    os.makedirs(out_dir)\n    tool = shutil.which(\"7z\") or shutil.which(\"7za\") or shutil.which(\"bsdtar\")\n    if tool and tool.endswith(\"bsdtar\"):\n        subprocess.run([tool, \"-xf\", archive_path, \"-C\", out_dir], check=True)\n    elif tool:\n        subprocess.run([tool, \"x\", \"-y\", f\"-o{out_dir}\", archive_path], check=True, stdout=subprocess.DEVNULL)\n    else:\n        pip_install(\"py7zr\")\n        import py7zr\n        with py7zr.SevenZipFile(archive_path, \"r\") as archive:\n            archive.extractall(path=out_dir)\n    return out_dir\n\ndef decode_png(data):\n    return np.asarray(Image.open(io.BytesIO(data)).convert(\"RGB\"), dtype=np.uint8)\n\ndef image_id(path):\n    return int(os.path.splitext(os.path.basename(path))[0])      # \"test/123.png\" -> 123\n\ndef read_png_archive(archive_path, n_images):\n    # Returns a uint8 array (n_images, 32, 32, 3) in which row k holds the image with id k+1.\n    images = np.zeros((n_images, 32, 32, 3), dtype=np.uint8)\n    found = np.zeros(n_images, dtype=bool)\n    start = time.time()\n\n    libarchive = load_libarchive()\n    if libarchive is not None:\n        print(f\"Plan A: streaming {os.path.basename(archive_path)} with libarchive\", flush=True)\n        count = 0\n        with libarchive.file_reader(archive_path) as archive:\n            for entry in archive:\n                if not (entry.isfile and entry.pathname.endswith(\".png\")):\n                    continue\n                k = image_id(entry.pathname) - 1\n                images[k] = decode_png(b\"\".join(entry.get_blocks()))\n                found[k] = True\n                count += 1\n                if count % 50000 == 0:\n                    print(f\"  {count:,} images decoded ({time.time() - start:.0f}s)\", flush=True)\n    else:\n        print(f\"Plan B: extracting {os.path.basename(archive_path)} to a folder\", flush=True)\n        out_dir = extract_to_folder(archive_path)\n        files = glob.glob(os.path.join(out_dir, \"**\", \"*.png\"), recursive=True)\n        print(f\"  files counted in the extracted folder: {len(files):,}\", flush=True)\n        for path in files:\n            k = image_id(path) - 1\n            with open(path, \"rb\") as f:\n                images[k] = decode_png(f.read())\n            found[k] = True\n        shutil.rmtree(out_dir, ignore_errors=True)\n\n    n_found = int(found.sum())\n    print(f\"Decoded {n_found:,} of {n_images:,} images in {time.time() - start:.0f}s\", flush=True)\n    if n_found != n_images:\n        raise RuntimeError(\"Some images are missing, so the archive was not read completely.\")\n    return images","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- Step 1.4: read the 300,000 Kaggle test images (cached in /tmp in case the cell is run again) ----\nTEST_CACHE = \"/tmp/kaggle_test_uint8.npy\"\nif os.path.exists(TEST_CACHE):\n    X_kaggle = np.load(TEST_CACHE)\n    print(\"Loaded the cached test array\", flush=True)\nelse:\n    X_kaggle = read_png_archive(TEST_7Z, N_TEST)\n    np.save(TEST_CACHE, X_kaggle)\n\nprint(\"Kaggle test array:\", X_kaggle.shape, X_kaggle.dtype, f\"({X_kaggle.nbytes / 1e6:.0f} MB)\", flush=True)\n\nfig, axes = plt.subplots(1, 8, figsize=(12, 1.8))\nfor ax, k in zip(axes, range(8)):\n    ax.imshow(X_kaggle[k])\n    ax.set_title(f\"id {k + 1}\", fontsize=9)\n    ax.axis(\"off\")\nplt.suptitle(\"First Kaggle test images (check that they decoded correctly)\", fontsize=10)\nplt.show()","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 2: Data loading and augmentation\n\nThe official CIFAR-10 training set (50,000 images) is used for training and the official CIFAR-10 test set (10,000 images) is used for validation.\nThe 10,000 official test images are the images Kaggle scores, so the validation accuracy is a close estimate of the Kaggle score.\nIf the download fails (Internet off), the code falls back to the competition's `train.7z` and keeps 5,000 images aside for validation.\n\nBoth data sets are small (150 MB as uint8), so they are kept on the GPU and augmented there one whole batch at a time.\nKaggle gives only 4 CPU cores, and per-image CPU transforms would leave the GPU waiting.\n\nAugmentations used for training:\n1. **Random crop**: pad 4 pixels (reflect) and cut a random 32×32 window, the same as `RandomCrop(32, padding=4)`.\n2. **Random horizontal flip** with probability 0.5.\n3. **CutMix**: with probability 0.5 per batch, a random box from another image in the batch is pasted in, and the label is mixed by box area.","metadata":{}},{"cell_type":"code","source":"# ---- Step 2.1: load the official CIFAR-10 training and test sets ----\nCLASSES = [\"airplane\", \"automobile\", \"bird\", \"cat\", \"deer\", \"dog\", \"frog\", \"horse\", \"ship\", \"truck\"]\n\ntry:\n    train_set = torchvision.datasets.CIFAR10(\"/tmp/cifar10\", train=True, download=True)\n    test_set = torchvision.datasets.CIFAR10(\"/tmp/cifar10\", train=False, download=True)\n    assert list(train_set.classes) == CLASSES\n    X_train, y_train = train_set.data, np.array(train_set.targets)\n    X_val, y_val = test_set.data, np.array(test_set.targets)\n    VAL_NAME = \"official CIFAR-10 test set\"\nexcept Exception as error:\n    print(\"CIFAR-10 download failed (is Internet on?):\", error, flush=True)\n    print(\"Using the competition train.7z instead, with a 5,000 image hold-out for validation.\", flush=True)\n    labels = pd.read_csv(find_file(\"trainLabels.csv\")).sort_values(\"id\")\n    assert (labels[\"id\"].values == np.arange(1, len(labels) + 1)).all()\n    X_all = read_png_archive(find_file(\"train.7z\"), len(labels))\n    y_all = labels[\"label\"].map({name: i for i, name in enumerate(CLASSES)}).to_numpy()\n    order = np.random.RandomState(SEED).permutation(len(y_all))\n    n_hold = min(5000, len(y_all) // 10)\n    X_train, y_train = X_all[order[n_hold:]], y_all[order[n_hold:]]\n    X_val, y_val = X_all[order[:n_hold]], y_all[order[:n_hold]]\n    VAL_NAME = \"5,000 image hold-out from train.7z\"\n\nprint(f\"Training images: {X_train.shape} | validation images ({VAL_NAME}): {X_val.shape}\", flush=True)\nprint(\"Images per class in training:\", np.bincount(y_train, minlength=10), flush=True)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- Step 2.2: put the data on the GPU and define the augmentations ----\ndef to_uint8_tensor(images):\n    # (N, 32, 32, 3) uint8 NumPy array -> (N, 3, 32, 32) uint8 tensor\n    return torch.from_numpy(images).permute(0, 3, 1, 2).contiguous()\n\nXtr = to_uint8_tensor(X_train).to(device)\nytr = torch.from_numpy(y_train).long().to(device)\nXva = to_uint8_tensor(X_val).to(device)\nyva = torch.from_numpy(y_val).long().to(device)\n\n# CIFAR-10 channel mean and standard deviation, on the 0-255 pixel scale\nMEAN = torch.tensor([0.4914, 0.4822, 0.4465], device=device).view(1, 3, 1, 1) * 255\nSTD = torch.tensor([0.2470, 0.2435, 0.2616], device=device).view(1, 3, 1, 1) * 255\n\ndef normalize(x_uint8):\n    return (x_uint8.float() - MEAN) / STD\n\ndef random_crop(x, pad=4):\n    # Same as RandomCrop(32, padding=4) but for a whole batch: pad, then cut a 32x32 window at a random offset per image.\n    n, c, h, w = x.shape\n    x = F.pad(x, (pad, pad, pad, pad), mode=\"reflect\")\n    dy = torch.randint(0, 2 * pad + 1, (n,), device=x.device)\n    dx = torch.randint(0, 2 * pad + 1, (n,), device=x.device)\n    rows = (dy[:, None] + torch.arange(h, device=x.device))[:, None, :, None]    # (n, 1, h, 1)\n    cols = (dx[:, None] + torch.arange(w, device=x.device))[:, None, None, :]    # (n, 1, 1, w)\n    n_idx = torch.arange(n, device=x.device)[:, None, None, None]\n    c_idx = torch.arange(c, device=x.device)[None, :, None, None]\n    return x[n_idx, c_idx, rows, cols]\n\ndef random_flip(x):\n    flip = torch.rand(x.size(0), device=x.device) < 0.5\n    return torch.where(flip[:, None, None, None], x.flip(3), x)\n\ndef augment(x_uint8):\n    return random_flip(random_crop(normalize(x_uint8)))\n\ndef cutmix(x, y, alpha=1.0):\n    # CutMix (Yun et al., 2019). The box area follows lam ~ Beta(alpha, alpha); every image in the batch receives a\n    # patch from another image (random pairing), and the target becomes lam * y_a + (1 - lam) * y_b.\n    lam = np.random.beta(alpha, alpha)\n    partner = torch.randperm(x.size(0), device=x.device)\n    h, w = x.shape[2], x.shape[3]\n    cut_h, cut_w = int(h * math.sqrt(1 - lam)), int(w * math.sqrt(1 - lam))\n    cy, cx = int(np.random.randint(h)), int(np.random.randint(w))\n    y1, y2 = max(cy - cut_h // 2, 0), min(cy + cut_h // 2, h)\n    x1, x2 = max(cx - cut_w // 2, 0), min(cx + cut_w // 2, w)\n    x[:, :, y1:y2, x1:x2] = x[partner, :, y1:y2, x1:x2]\n    lam = 1 - (y2 - y1) * (x2 - x1) / (h * w)      # recompute from the real box, since it can be clipped at the border\n    return x, y, y[partner], lam\n\n# Show what the network actually sees during training\ndemo, _, _, demo_lam = cutmix(augment(Xtr[:16]), ytr[:16])\ndemo = (demo * STD + MEAN).clamp(0, 255).byte().permute(0, 2, 3, 1).cpu().numpy()\nfig, axes = plt.subplots(2, 8, figsize=(12, 3.4))\nfor ax, img, label in zip(axes.flat, demo, y_train[:16]):\n    ax.imshow(img)\n    ax.set_title(CLASSES[label], fontsize=8)\n    ax.axis(\"off\")\nplt.suptitle(f\"Random crop + flip + CutMix (each image keeps {demo_lam:.0%} of its own pixels)\", fontsize=10)\nplt.tight_layout()\nplt.savefig(f\"{WORK_DIR}/augmentation_examples.png\", dpi=150)\nplt.show()","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- Step 2.3: Dataset and DataLoader for the 300,000 Kaggle test images ----\nclass KaggleTestDataset(Dataset):\n    # Returns the test images in id order (index k = id k+1) as uint8 tensors; normalisation is done on the GPU.\n    def __init__(self, images):\n        self.images = images\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, index):\n        return torch.from_numpy(self.images[index]).permute(2, 0, 1)\n\ntest_loader = DataLoader(KaggleTestDataset(X_kaggle), batch_size=1000, shuffle=False,\n                         num_workers=2, pin_memory=True)\nprint(f\"Kaggle test loader: {len(test_loader.dataset):,} images in {len(test_loader)} batches\", flush=True)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 3: Custom network design\n\nThe network is a pre-activation residual network with Squeeze-and-Excitation (SE) attention in every block. It is built only from\n`nn.Conv2d`, `nn.BatchNorm2d`, `nn.Linear` and `nn.Dropout`.\n\n| Stage | Output size | Blocks | Channels |\n|---|---|---|---|\n| Stem: 3×3 conv | 32×32 | – | 64 |\n| Stage 1 | 32×32 | 2 | 64 |\n| Stage 2 (stride 2) | 16×16 | 2 | 128 |\n| Stage 3 (stride 2) | 8×8 | 2 | 256 |\n| Stage 4 (stride 2) | 4×4 | 2 | 512 |\n| Head: BN, GELU, global average pool, dropout 0.1, linear | 1×1 | – | 10 |\n\nEach block: BN → GELU → 3×3 conv → BN → GELU → 3×3 conv → SE → add the shortcut. When the size or width changes,\nthe shortcut is a 1×1 convolution with stride 2.","metadata":{}},{"cell_type":"code","source":"# ---- Step 3: the custom network ----\n\nclass SqueezeExcite(nn.Module):\n    def __init__(self, channels, reduction=16):\n        super().__init__()\n        hidden = max(channels // reduction, 8)\n        self.fc1 = nn.Conv2d(channels, hidden, kernel_size=1)\n        self.fc2 = nn.Conv2d(hidden, channels, kernel_size=1)\n\n    def forward(self, x):\n        w = F.adaptive_avg_pool2d(x, 1)                     # squeeze: one number per channel\n        w = torch.sigmoid(self.fc2(F.gelu(self.fc1(w))))    # excitation: channel weights between 0 and 1\n        return x * w\n\nclass PreActSEBlock(nn.Module):\n    def __init__(self, in_ch, out_ch, stride=1):\n        super().__init__()\n        self.bn1 = nn.BatchNorm2d(in_ch)\n        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, stride=stride, padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(out_ch)\n        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, stride=1, padding=1, bias=False)\n        self.se = SqueezeExcite(out_ch)\n        self.shortcut = None\n        if stride != 1 or in_ch != out_ch:\n            self.shortcut = nn.Conv2d(in_ch, out_ch, 1, stride=stride, bias=False)\n\n    def forward(self, x):\n        out = F.gelu(self.bn1(x))\n        skip = x if self.shortcut is None else self.shortcut(out)\n        out = self.conv1(out)\n        out = self.conv2(F.gelu(self.bn2(out)))\n        out = self.se(out)\n        return out + skip\n\nclass CustomSEResNet(nn.Module):\n    def __init__(self, num_classes=10, widths=(64, 128, 256, 512), blocks_per_stage=(2, 2, 2, 2), dropout=0.1):\n        super().__init__()\n        self.stem = nn.Conv2d(3, widths[0], 3, stride=1, padding=1, bias=False)\n        blocks, in_ch = [], widths[0]\n        for stage, (width, n_blocks) in enumerate(zip(widths, blocks_per_stage)):\n            for b in range(n_blocks):\n                stride = 2 if (stage > 0 and b == 0) else 1\n                blocks.append(PreActSEBlock(in_ch, width, stride))\n                in_ch = width\n        self.blocks = nn.Sequential(*blocks)\n        self.bn_final = nn.BatchNorm2d(in_ch)\n        self.dropout = nn.Dropout(dropout)\n        self.fc = nn.Linear(in_ch, num_classes)\n\n        # Random initialisation (no pre-trained weights): Kaiming normal for convolutions, BN starts as identity.\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                if m.bias is not None:\n                    nn.init.zeros_(m.bias)\n            elif isinstance(m, nn.BatchNorm2d):\n                nn.init.ones_(m.weight)\n                nn.init.zeros_(m.bias)\n\n    def forward(self, x):\n        x = self.stem(x)\n        x = self.blocks(x)\n        x = F.gelu(self.bn_final(x))\n        x = F.adaptive_avg_pool2d(x, 1).flatten(1)\n        return self.fc(self.dropout(x))\n\nmodel = CustomSEResNet().to(device, memory_format=torch.channels_last)\nn_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n\n# Layer-by-layer summary (output shape and parameters), using one dummy image\ndef count(module):\n    return sum(p.numel() for p in module.parameters())\n\nrows = []\nmodel.eval()\nwith torch.no_grad():\n    x = model.stem(torch.zeros(1, 3, 32, 32, device=device))\n    rows.append((\"stem conv 3x3\", tuple(x.shape[1:]), count(model.stem)))\n    for i, block in enumerate(model.blocks):\n        x = block(x)\n        rows.append((f\"PreActSEBlock {i + 1}\", tuple(x.shape[1:]), count(block)))\n    rows.append((\"head (BN, GELU, pool, dropout, linear)\", (10,), count(model.bn_final) + count(model.fc)))\nmodel.train()\nprint(pd.DataFrame(rows, columns=[\"layer\", \"output shape\", \"parameters\"]).to_string(index=False))\nprint(f\"\\nTotal trainable parameters: {n_params:,}\", flush=True)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 4: Training with AdamW, warmup + cosine annealing and CutMix\n\nRecommended settings (used below): 120 epochs, batch size 512, peak learning rate 2e-3, weight decay 0.05, 5 warmup epochs.\nCutMix is applied to half of the batches and is switched off for the last 10 epochs, so the network finishes on clean images.\nMixed precision (float16) is used for speed, and both T4 GPUs are used through `nn.DataParallel`.\n\nWith CutMix the target of a mixed batch is `lam * y_a + (1 - lam) * y_b`, so the loss is computed as\n`lam * CE(logits, y_a) + (1 - lam) * CE(logits, y_b)`.\nBecause of this, the training loss stays higher than the validation loss; that is expected.","metadata":{}},{"cell_type":"code","source":"# ---- Step 4.1: hyperparameters, optimizer and scheduler ----\nEPOCHS = 120\nBATCH_SIZE = 512\nPEAK_LR = 2e-3\nWEIGHT_DECAY = 0.05\nWARMUP_EPOCHS = 5\nCUTMIX_PROB = 0.5          # share of batches that get CutMix\nCUTMIX_ALPHA = 1.0         # Beta(1, 1): box size is uniform\nCLEAN_LAST_EPOCHS = 10     # no CutMix in the final epochs\n\nsteps_per_epoch = len(Xtr) // BATCH_SIZE\ntotal_steps = EPOCHS * steps_per_epoch\nwarmup_steps = WARMUP_EPOCHS * steps_per_epoch\n\n# AdamW: decay only the convolution and linear weights. BatchNorm scales/shifts and biases are left without decay.\ndecay = [p for p in model.parameters() if p.ndim > 1]\nno_decay = [p for p in model.parameters() if p.ndim <= 1]\noptimizer = torch.optim.AdamW([{\"params\": decay, \"weight_decay\": WEIGHT_DECAY},\n                               {\"params\": no_decay, \"weight_decay\": 0.0}],\n                              lr=PEAK_LR, betas=(0.9, 0.999), eps=1e-8)\n\n# Linear warmup from almost 0 to PEAK_LR, then cosine annealing down to 0. Updated after every batch.\ndef warmup_cosine(step):\n    if step < warmup_steps:\n        return (step + 1) / warmup_steps\n    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n    return 0.5 * (1 + math.cos(math.pi * progress))\n\nscheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=warmup_cosine)\ncriterion = nn.CrossEntropyLoss()\n\nuse_amp = True\ntry:\n    scaler = torch.amp.GradScaler(\"cuda\", enabled=use_amp)\nexcept (AttributeError, TypeError):\n    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)\n\n# Use both GPUs when Kaggle gives two (GPU T4 x2). \"model\" stays the plain network for saving.\nnet = nn.DataParallel(model) if n_gpus > 1 else model\n\nprint(f\"{EPOCHS} epochs x {steps_per_epoch} steps = {total_steps:,} steps, warmup {warmup_steps} steps, \"\n      f\"{max(n_gpus, 1)} GPU(s)\", flush=True)\n\n@torch.no_grad()\ndef predict_probs(x_uint8, tta=False):\n    # Class probabilities for a batch of uint8 images; with tta=True the horizontally flipped copy is averaged in.\n    x = normalize(x_uint8).contiguous(memory_format=torch.channels_last)\n    with torch.autocast(device_type=\"cuda\", dtype=torch.float16, enabled=use_amp):\n        probs = net(x).float().softmax(1)\n        if tta:\n            probs = 0.5 * (probs + net(x.flip(3)).float().softmax(1))\n    return probs\n\n@torch.no_grad()\ndef evaluate(X, y, tta=False, batch_size=1000):\n    net.eval()\n    loss_sum, correct = 0.0, 0\n    for i in range(0, len(X), batch_size):\n        probs = predict_probs(X[i:i + batch_size], tta)\n        target = y[i:i + batch_size]\n        loss_sum += F.nll_loss(probs.clamp_min(1e-8).log(), target, reduction=\"sum\").item()\n        correct += (probs.argmax(1) == target).sum().item()\n    return loss_sum / len(X), correct / len(X)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- Step 4.2: training loop ----\nhistory, lr_trace = [], []\ntrain_start = time.time()\n\nfor epoch in range(1, EPOCHS + 1):\n    net.train()\n    use_cutmix = epoch <= EPOCHS - CLEAN_LAST_EPOCHS\n    order = torch.randperm(len(Xtr), device=device)\n    loss_sum = torch.zeros((), device=device)\n    epoch_start = time.time()\n\n    for step in range(steps_per_epoch):\n        idx = order[step * BATCH_SIZE:(step + 1) * BATCH_SIZE]\n        x, y = augment(Xtr[idx]), ytr[idx]\n        mixed = use_cutmix and np.random.rand() < CUTMIX_PROB\n        if mixed:\n            x, y_a, y_b, lam = cutmix(x, y, CUTMIX_ALPHA)\n        x = x.contiguous(memory_format=torch.channels_last)\n\n        with torch.autocast(device_type=\"cuda\", dtype=torch.float16, enabled=use_amp):\n            logits = net(x)\n            if mixed:\n                loss = lam * criterion(logits, y_a) + (1 - lam) * criterion(logits, y_b)\n            else:\n                loss = criterion(logits, y)\n\n        optimizer.zero_grad(set_to_none=True)\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        scheduler.step()\n\n        loss_sum += loss.detach()\n        lr_trace.append(scheduler.get_last_lr()[0])\n\n    train_loss = (loss_sum / steps_per_epoch).item()\n    val_loss, val_acc = evaluate(Xva, yva)\n    seconds = time.time() - epoch_start\n    lr_now = optimizer.param_groups[0][\"lr\"]\n    history.append({\"epoch\": epoch, \"lr\": lr_now, \"train_loss\": train_loss, \"val_loss\": val_loss,\n                    \"val_acc\": val_acc, \"cutmix\": use_cutmix, \"seconds\": seconds})\n    print(f\"Epoch {epoch:3d}/{EPOCHS} | lr {lr_now:.2e} | train loss {train_loss:.4f} | \"\n          f\"val loss {val_loss:.4f} | val acc {val_acc * 100:.2f}% | {seconds:.0f}s\", flush=True)\n\ntrain_minutes = (time.time() - train_start) / 60\ntorch.save(model.state_dict(), f\"{WORK_DIR}/custom_se_resnet_cifar10.pth\")\nlog = pd.DataFrame(history)\nlog.to_csv(f\"{WORK_DIR}/training_log.csv\", index=False)\nprint(f\"\\nTraining finished in {train_minutes:.1f} minutes\", flush=True)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- Step 4.3: final evaluation and training curves ----\nfinal_loss, final_acc = evaluate(Xva, yva, tta=False)\ntta_loss, tta_acc = evaluate(Xva, yva, tta=True)\nprint(f\"Final accuracy on the {VAL_NAME}: {final_acc * 100:.2f}% (with flip TTA: {tta_acc * 100:.2f}%)\", flush=True)\n\nfig, axes = plt.subplots(1, 3, figsize=(16, 4.2))\naxes[0].plot(log.epoch, log.train_loss, label=\"train loss (CutMix targets)\")\naxes[0].plot(log.epoch, log.val_loss, label=\"validation loss\")\naxes[0].set_xlabel(\"epoch\"); axes[0].set_ylabel(\"cross-entropy\"); axes[0].set_title(\"Loss\"); axes[0].legend()\n\naxes[1].plot(log.epoch, log.val_acc * 100, label=\"validation accuracy\")\naxes[1].axhline(94.2, color=\"red\", linestyle=\"--\", label=\"target 94.2%\")\naxes[1].set_ylim(max(0, (log.val_acc.min() * 100) - 2), 100)\naxes[1].set_xlabel(\"epoch\"); axes[1].set_ylabel(\"accuracy (%)\"); axes[1].set_title(\"Validation accuracy\"); axes[1].legend()\n\naxes[2].plot(np.arange(1, len(lr_trace) + 1) / steps_per_epoch, lr_trace)\naxes[2].set_xlabel(\"epoch\"); axes[2].set_ylabel(\"learning rate\"); axes[2].set_title(\"Warmup + cosine annealing\")\n\nplt.tight_layout()\nplt.savefig(f\"{WORK_DIR}/training_curves.png\", dpi=150)\nplt.show()","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 5: Inference on the 300,000 Kaggle test images and submission file\n\nPredictions use test-time augmentation: the softmax of each image and of its mirror image are averaged.\nClass indices are mapped back to the Kaggle label names and written to `submission.csv` in the same id order as `sampleSubmission.csv`.","metadata":{}},{"cell_type":"code","source":"# ---- Step 5: predict the Kaggle test set and write submission.csv ----\nnet.eval()\nstart = time.time()\npredictions = []\nfor xb in test_loader:\n    probs = predict_probs(xb.to(device, non_blocking=True), tta=True)\n    predictions.append(probs.argmax(1).cpu())\npredictions = torch.cat(predictions).numpy()\nprint(f\"Predicted {len(predictions):,} images in {time.time() - start:.0f}s\", flush=True)\n\nsample = pd.read_csv(SAMPLE_SUB)\nsubmission = pd.DataFrame({\"id\": sample[\"id\"].values,\n                           \"label\": np.array(CLASSES)[predictions[sample[\"id\"].values - 1]]})\nsubmission.to_csv(f\"{WORK_DIR}/submission.csv\", index=False)\n\nassert len(submission) == len(sample) and submission[\"label\"].isin(CLASSES).all()\nprint(\"submission.csv saved:\", submission.shape, flush=True)\nprint(submission.head(10).to_string(index=False))\nprint(\"\\nPredicted label counts:\\n\" + submission[\"label\"].value_counts().to_string(), flush=True)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 6: Summary of results for the report","metadata":{}},{"cell_type":"code","source":"# ---- Step 6: numbers needed for the report (also saved to results_summary.txt) ----\nbest = log.loc[log.val_acc.idxmax()]\nsummary = f'''Model: CustomSEResNet (custom pre-activation residual network with SE blocks, written from scratch)\nTrainable parameters: {n_params:,}\nInitialisation: random (Kaiming normal), no pre-trained weights\nOptimizer: AdamW, peak lr {PEAK_LR}, betas (0.9, 0.999), weight decay {WEIGHT_DECAY} (none on BatchNorm/bias)\nLR schedule: linear warmup for {WARMUP_EPOCHS} epochs, then cosine annealing to 0 ({total_steps:,} steps in total)\nAugmentation: random crop (4 px reflect padding) + horizontal flip + CutMix (p={CUTMIX_PROB}, alpha={CUTMIX_ALPHA}); CutMix off for the last {CLEAN_LAST_EPOCHS} epochs\nEpochs: {EPOCHS}, batch size: {BATCH_SIZE}, mixed precision: {use_amp}, GPUs: {max(n_gpus, 1)} x {torch.cuda.get_device_name(0)}\nTraining time: {train_minutes:.1f} minutes ({log.seconds.mean():.1f} s per epoch)\nValidation set: {VAL_NAME}\nFinal validation accuracy: {final_acc * 100:.2f}%\nFinal validation accuracy with flip TTA: {tta_acc * 100:.2f}%\nHighest validation accuracy during training: {best.val_acc * 100:.2f}% (epoch {int(best.epoch)})\nKaggle test images predicted: {len(submission):,}\n'''\nwith open(f\"{WORK_DIR}/results_summary.txt\", \"w\") as f:\n    f.write(summary)\nprint(summary)","metadata":{},"outputs":[],"execution_count":null}]}