{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.13.15","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","isInternetEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"c0","cell_type":"markdown","source":"# CIFAR-10 — RaySENet (custom network) + AdamW \nAssignment 2 of the CIFAR-10 Kaggle task.\n","metadata":{}},{"id":"c1","cell_type":"code","source":"!pip -q install py7zr","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T09:29:29.617387Z","iopub.execute_input":"2026-10-10T09:29:29.618046Z","iopub.status.idle":"2026-10-10T09:29:35.82875Z","shell.execute_reply.started":"2026-10-10T09:29:29.618015Z","shell.execute_reply":"2026-10-10T09:29:35.827994Z"}},"outputs":[],"execution_count":null},{"id":"c2","cell_type":"code","source":"import os, random, glob, time, math\nimport numpy as np\nimport torch, torch.nn as nn, torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision\nfrom PIL import Image\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nSEED = 42\nrandom.seed(SEED); np.random.seed(SEED)\ntorch.manual_seed(SEED); torch.cuda.manual_seed_all(SEED)\ntorch.backends.cudnn.benchmark = True\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint('Device:', device, torch.cuda.get_device_name(0) if device.type == 'cuda' else '')\n\nEPOCHS        = 120\nWARMUP_EPOCHS = 5\nBATCH_SIZE    = 256\nPEAK_LR       = 2e-3\nMIN_LR        = 1e-6\nWEIGHT_DECAY  = 0.05\nLABEL_SMOOTH  = 0.1\nCUTMIX_PROB   = 0.5\nCUTOUT_SIZE   = 16\n\nCLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer',\n           'dog', 'frog', 'horse', 'ship', 'truck']\nMEAN = torch.tensor([0.4914, 0.4822, 0.4465]).view(1, 3, 1, 1)\nSTD  = torch.tensor([0.2470, 0.2435, 0.2616]).view(1, 3, 1, 1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T09:29:35.830685Z","iopub.execute_input":"2026-10-10T09:29:35.830956Z","iopub.status.idle":"2026-10-10T09:29:47.903424Z","shell.execute_reply.started":"2026-10-10T09:29:35.830933Z","shell.execute_reply":"2026-10-10T09:29:47.902687Z"}},"outputs":[],"execution_count":null},{"id":"c3","cell_type":"markdown","source":"## 1. Data\nThe Kaggle training set is identical to the official CIFAR-10 training set, so we load it with torchvision and keep the whole thing on the GPU as a normalised tensor. The official 10,000-image test set is used only to *estimate* the Kaggle score — it is never trained on.","metadata":{}},{"id":"c4","cell_type":"code","source":"def to_tensor(ds):\n    x = torch.tensor(ds.data).permute(0, 3, 1, 2).float().div(255)   # N,3,32,32 in [0,1]\n    x = (x - MEAN) / STD\n    y = torch.tensor(ds.targets, dtype=torch.long)\n    return x, y\n\ntrain_raw = torchvision.datasets.CIFAR10('./data', train=True,  download=True)\ntest_raw  = torchvision.datasets.CIFAR10('./data', train=False, download=True)\nX_train, Y_train = [t.to(device) for t in to_tensor(train_raw)]\nX_test,  Y_test  = [t.to(device) for t in to_tensor(test_raw)]\nprint('train', X_train.shape, ' test', X_test.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T09:29:47.904325Z","iopub.execute_input":"2026-10-10T09:29:47.904795Z","iopub.status.idle":"2026-10-10T09:57:01.787299Z","shell.execute_reply.started":"2026-10-10T09:29:47.904773Z","shell.execute_reply":"2026-10-10T09:57:01.786575Z"}},"outputs":[],"execution_count":null},{"id":"c5","cell_type":"markdown","source":"## 2. Augmentation (all on GPU)\n* **Random crop**  and **horizontal flip**.\n* **Cutout**: a random 16×16 square of every image is set to 0.\n* **CutMix**: with probability 0.5 per batch, a random box from a shuffled copy of the batch is pasted in; the label becomes a mix of the two classes weighted by box area, λ ~ Beta(1, 1).","metadata":{}},{"id":"c6","cell_type":"code","source":"def random_crop_flip(x, pad=4):\n    B, C, H, W = x.shape\n    xp = F.pad(x, (pad, pad, pad, pad), 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))[:, None, :, None]\n    cols = (j[:, None] + torch.arange(W, device=x.device))[:, None, None, :]\n    bidx = torch.arange(B, device=x.device)[:, None, None, None]\n    cidx = torch.arange(C, device=x.device)[None, :, None, None]\n    out = xp[bidx, cidx, rows, cols]\n    flip = torch.rand(B, device=x.device) < 0.5\n    return torch.where(flip[:, None, None, None], out.flip(3), out)\n\ndef cutout(x, size=CUTOUT_SIZE):\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)[None, :, None]\n    xs = torch.arange(W, device=x.device)[None, None, :]\n    h = size // 2\n    hole = ((ys - cy) >= -h) & ((ys - cy) < h) & ((xs - cx) >= -h) & ((xs - cx) < h)   # size x size square\n    return x * (~hole)[:, None].float()\n\ndef cutmix(x, y):\n    B, _, H, W = x.shape\n    lam = np.random.beta(1.0, 1.0)\n    perm = torch.randperm(B, device=x.device)\n    rh, rw = int(H * math.sqrt(1 - lam)), int(W * math.sqrt(1 - lam))\n    cy, cx = np.random.randint(H), np.random.randint(W)\n    y1, y2 = max(cy - rh // 2, 0), min(cy + rh // 2, H)\n    x1, x2 = max(cx - rw // 2, 0), min(cx + rw // 2, W)\n    x = x.clone()\n    x[:, :, y1:y2, x1:x2] = x[perm, :, y1:y2, x1:x2]\n    lam = 1 - (y2 - y1) * (x2 - x1) / (H * W)       # exact area ratio\n    return x, y, y[perm], lam","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T09:57:01.788446Z","iopub.execute_input":"2026-10-10T09:57:01.788774Z","iopub.status.idle":"2026-10-10T09:57:01.799092Z","shell.execute_reply.started":"2026-10-10T09:57:01.788753Z","shell.execute_reply":"2026-10-10T09:57:01.798563Z"}},"outputs":[],"execution_count":null},{"id":"c7","cell_type":"markdown","source":"## 3. RaySENet — the custom network\n","metadata":{}},{"id":"c8","cell_type":"code","source":"class SE(nn.Module):\n    def __init__(self, c, r=8):\n        super().__init__()\n        self.fc1 = nn.Conv2d(c, c // r, 1)\n        self.fc2 = nn.Conv2d(c // r, c, 1)\n    def forward(self, x):\n        s = F.adaptive_avg_pool2d(x, 1)\n        s = torch.sigmoid(self.fc2(F.relu(self.fc1(s), inplace=True)))\n        return x * s\n\nclass SEBlock(nn.Module):\n    def __init__(self, cin, cout, stride=1):\n        super().__init__()\n        self.conv1 = nn.Conv2d(cin, cout, 3, stride, 1, bias=False)\n        self.bn1   = nn.BatchNorm2d(cout)\n        self.conv2 = nn.Conv2d(cout, cout, 3, 1, 1, bias=False)\n        self.bn2   = nn.BatchNorm2d(cout)\n        nn.init.zeros_(self.bn2.weight)                      # zero-init residual gate\n        self.se    = SE(cout)\n        if stride != 1 or cin != cout:\n            self.short = nn.Sequential(\n                nn.AvgPool2d(stride) if stride > 1 else nn.Identity(),\n                nn.Conv2d(cin, cout, 1, bias=False), nn.BatchNorm2d(cout))\n        else:\n            self.short = nn.Identity()\n    def forward(self, x):\n        out = F.relu(self.bn1(self.conv1(x)), inplace=True)\n        out = self.se(self.bn2(self.conv2(out)))\n        return F.relu(out + self.short(x), inplace=True)\n\nclass RaySENet(nn.Module):\n    def __init__(self, num_classes=10, widths=(64, 128, 256, 512), depths=(2, 2, 3, 2), drop=0.2):\n        super().__init__()\n        self.stem = nn.Sequential(\n            nn.Conv2d(3, 32, 3, 1, 1, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True),\n            nn.Conv2d(32, widths[0], 3, 1, 1, bias=False), nn.BatchNorm2d(widths[0]), nn.ReLU(inplace=True))\n        layers, cin = [], widths[0]\n        for s, (w, d) in enumerate(zip(widths, depths)):\n            for b in range(d):\n                layers.append(SEBlock(cin, w, stride=2 if (b == 0 and s > 0) else 1))\n                cin = w\n        self.stages = nn.Sequential(*layers)\n        self.drop = nn.Dropout(drop)\n        self.fc = nn.Linear(2 * cin, num_classes)\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n                if m.bias is not None: nn.init.zeros_(m.bias)\n    def forward(self, x):\n        x = self.stages(self.stem(x))\n        x = torch.cat([F.adaptive_avg_pool2d(x, 1), F.adaptive_max_pool2d(x, 1)], 1).flatten(1)\n        return self.fc(self.drop(x))\n\nmodel = RaySENet().to(device).to(memory_format=torch.channels_last)\nprint(f'RaySENet parameters: {sum(p.numel() for p in model.parameters())/1e6:.2f} M')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T09:57:01.801133Z","iopub.execute_input":"2026-10-10T09:57:01.8016Z","iopub.status.idle":"2026-10-10T09:57:02.043414Z","shell.execute_reply.started":"2026-10-10T09:57:01.801574Z","shell.execute_reply":"2026-10-10T09:57:02.042608Z"}},"outputs":[],"execution_count":null},{"id":"c9","cell_type":"markdown","source":"## 4. Optimiser and learning-rate schedule\n","metadata":{}},{"id":"c10","cell_type":"code","source":"decay, no_decay = [], []\nfor n, p in model.named_parameters():\n    (no_decay if p.ndim <= 1 else decay).append(p)\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))\n\nsteps_per_epoch = math.ceil(len(X_train) / BATCH_SIZE)\ntotal_steps  = EPOCHS * steps_per_epoch\nwarmup_steps = WARMUP_EPOCHS * steps_per_epoch\n\ndef lr_factor(step):\n    if step < warmup_steps:\n        return (step + 1) / warmup_steps\n    t = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n    return MIN_LR / PEAK_LR + (1 - MIN_LR / PEAK_LR) * 0.5 * (1 + math.cos(math.pi * t))\n\nscheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_factor)\ncriterion = nn.CrossEntropyLoss(label_smoothing=LABEL_SMOOTH)\nuse_amp = device.type == 'cuda'\nscaler = torch.amp.GradScaler('cuda', enabled=use_amp)\n\nlrs = [PEAK_LR * lr_factor(s) for s in range(0, total_steps, 10)]\nplt.figure(figsize=(6, 2.5)); plt.plot(np.arange(len(lrs)) * 10 / steps_per_epoch, lrs)\nplt.xlabel('epoch'); plt.ylabel('learning rate'); plt.title('Warmup + cosine schedule'); plt.tight_layout(); plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T09:57:02.044489Z","iopub.execute_input":"2026-10-10T09:57:02.045284Z","iopub.status.idle":"2026-10-10T09:57:02.25297Z","shell.execute_reply.started":"2026-10-10T09:57:02.045221Z","shell.execute_reply":"2026-10-10T09:57:02.2524Z"}},"outputs":[],"execution_count":null},{"id":"c11","cell_type":"markdown","source":"## 5. Training","metadata":{}},{"id":"c12","cell_type":"code","source":"@torch.no_grad()\ndef predict_logits(x_all, bs=1000, tta=True):\n    model.eval()\n    outs = []\n    for k in range(0, len(x_all), bs):\n        xb = x_all[k:k + bs].to(device, non_blocking=True).contiguous(memory_format=torch.channels_last)\n        with torch.autocast('cuda', dtype=torch.float16, enabled=use_amp):\n            out = model(xb).float().softmax(1)\n            if tta:\n                out = out + model(xb.flip(3)).float().softmax(1)\n        outs.append(out)\n    return torch.cat(outs)\n\ndef evaluate(tta=False):\n    return (predict_logits(X_test, tta=tta).argmax(1) == Y_test).float().mean().item()\n\ndef train_one_epoch():\n    model.train()\n    perm = torch.randperm(len(X_train), device=device)\n    tot_loss, n = 0.0, 0\n    for k in range(0, len(X_train), BATCH_SIZE):\n        idx = perm[k:k + BATCH_SIZE]\n        xb, yb = X_train[idx], Y_train[idx]\n        xb = cutout(random_crop_flip(xb))\n        mixed = np.random.rand() < CUTMIX_PROB\n        if mixed:\n            xb, ya, yb2, lam = cutmix(xb, yb)\n        xb = xb.contiguous(memory_format=torch.channels_last)\n        with torch.autocast('cuda', dtype=torch.float16, enabled=use_amp):\n            out = model(xb)\n            loss = (lam * criterion(out, ya) + (1 - lam) * criterion(out, yb2)) if mixed else criterion(out, yb)\n        optimizer.zero_grad(set_to_none=True)\n        scaler.scale(loss).backward()\n        scaler.unscale_(optimizer)\n        nn.utils.clip_grad_norm_(model.parameters(), 5.0)\n        scaler.step(optimizer); scaler.update(); scheduler.step()\n        tot_loss += loss.item() * len(idx); n += len(idx)\n    return tot_loss / n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T09:57:02.253974Z","iopub.execute_input":"2026-10-10T09:57:02.25464Z","iopub.status.idle":"2026-10-10T09:57:02.263287Z","shell.execute_reply.started":"2026-10-10T09:57:02.254616Z","shell.execute_reply":"2026-10-10T09:57:02.262679Z"}},"outputs":[],"execution_count":null},{"id":"c13","cell_type":"code","source":"history = {'epoch': [], 'loss': [], 'test_acc': [], 'lr': []}\nt0 = time.time()\nfor epoch in range(1, EPOCHS + 1):\n    loss = train_one_epoch()\n    if epoch % 5 == 0 or epoch == 1 or epoch > EPOCHS - 10:\n        acc = evaluate()\n        history['epoch'].append(epoch); history['loss'].append(loss)\n        history['test_acc'].append(acc); history['lr'].append(optimizer.param_groups[0]['lr'])\n        print(f'epoch {epoch:3d}/{EPOCHS} | loss {loss:.4f} | official-test acc {acc:.4f} | '\n              f'lr {optimizer.param_groups[0][\"lr\"]:.2e} | {(time.time() - t0) / 60:.1f} min')\ntorch.save(model.state_dict(), 'raysenet_cifar10.pt')\nprint(f'Training done in {(time.time() - t0) / 60:.1f} min')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T09:57:02.264339Z","iopub.execute_input":"2026-10-10T09:57:02.264646Z","iopub.status.idle":"2026-10-10T10:33:02.747145Z","shell.execute_reply.started":"2026-10-10T09:57:02.264614Z","shell.execute_reply":"2026-10-10T10:33:02.746217Z"}},"outputs":[],"execution_count":null},{"id":"c14","cell_type":"code","source":"acc_plain = evaluate(tta=False)\nacc_tta   = evaluate(tta=True)\nprint(f'Official CIFAR-10 test accuracy: {acc_plain:.4f}  |  with flip-TTA: {acc_tta:.4f}')\n\nfig, ax = plt.subplots(1, 2, figsize=(11, 3.5))\nax[0].plot(history['epoch'], history['loss']); ax[0].set_title('Training loss'); ax[0].set_xlabel('epoch')\nax[1].plot(history['epoch'], history['test_acc']); ax[1].set_title('Official test accuracy'); ax[1].set_xlabel('epoch')\nplt.tight_layout(); plt.savefig('training_curves.png', dpi=150); plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T10:33:02.748225Z","iopub.execute_input":"2026-10-10T10:33:02.748563Z","iopub.status.idle":"2026-10-10T10:33:06.334388Z","shell.execute_reply.started":"2026-10-10T10:33:02.74853Z","shell.execute_reply":"2026-10-10T10:33:06.333571Z"}},"outputs":[],"execution_count":null},{"id":"c15","cell_type":"markdown","source":"## 6. Predict the 300,000 Kaggle test images and write `submission.csv`","metadata":{}},{"id":"c16","cell_type":"code","source":"import py7zr\nTEST_7Z = glob.glob('/kaggle/input/**/test.7z', recursive=True)[0]\nprint('Using:', TEST_7Z)\nif not os.path.isdir('test') or len(glob.glob('test/*.png')) < 300000:\n    with py7zr.SevenZipFile(TEST_7Z, mode='r') as z:\n        z.extractall(path='.')\nprint('Images found:', len(glob.glob('test/*.png')))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T10:33:06.3354Z","iopub.execute_input":"2026-10-10T10:33:06.335702Z","iopub.status.idle":"2026-10-10T10:36:24.756549Z","shell.execute_reply.started":"2026-10-10T10:33:06.335674Z","shell.execute_reply":"2026-10-10T10:36:24.755684Z"}},"outputs":[],"execution_count":null},{"id":"c17","cell_type":"code","source":"class KaggleTest(Dataset):\n    def __init__(self, n=300000): self.ids = list(range(1, n + 1))\n    def __len__(self): return len(self.ids)\n    def __getitem__(self, i):\n        img = np.asarray(Image.open(f'test/{self.ids[i]}.png').convert('RGB'), dtype=np.uint8)\n        return torch.from_numpy(img.copy()).permute(2, 0, 1)\n\nloader = DataLoader(KaggleTest(), batch_size=1000, shuffle=False, num_workers=2, pin_memory=True)\npreds = []\nfor xb in loader:\n    xb = (xb.float().div(255) - MEAN) / STD\n    preds.append(predict_logits(xb, bs=1000, tta=True).argmax(1).cpu())\npreds = torch.cat(preds).numpy()\n\nsub = pd.DataFrame({'id': np.arange(1, 300001), 'label': [CLASSES[p] for p in preds]})\nsub.to_csv('submission.csv', index=False)\nprint(sub.head()); print(sub['label'].value_counts())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T10:36:24.757628Z","iopub.execute_input":"2026-10-10T10:36:24.757953Z","iopub.status.idle":"2026-10-10T10:37:31.781417Z","shell.execute_reply.started":"2026-10-10T10:36:24.757932Z","shell.execute_reply":"2026-10-10T10:37:31.780682Z"}},"outputs":[],"execution_count":null},{"id":"c18","cell_type":"markdown","source":"","metadata":{}}]}