{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# %% [Cell 1]\n# ===== 1. Setup & data loading =====\nimport os, glob, time, math, copy, random, shutil, subprocess\nimport numpy as np, pandas as pd\nfrom PIL import Image\nfrom concurrent.futures import ProcessPoolExecutor\nimport torch, torch.nn as nn, torch.nn.functional as F\n\nSEED = 42\nrandom.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)\ntorch.backends.cudnn.benchmark = True\nDEVICE = 'cuda'\nN_GPU = torch.cuda.device_count()\nprint('torch', torch.__version__, '| GPUs:', [torch.cuda.get_device_name(i) for i in range(N_GPU)])\n\nCLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']\nIN_DIR = os.path.dirname(glob.glob('/kaggle/input/**/trainLabels.csv', recursive=True)[0])\nWORK = '/kaggle/temp'                      # scratch space, not saved as notebook output\nos.makedirs(WORK, exist_ok=True)\nprint('competition files:', os.listdir(IN_DIR))\n\ndef extract(archive):\n    if shutil.which('7z'):\n        subprocess.run(['7z', 'x', '-y', '-o' + WORK, archive], check=True, stdout=subprocess.DEVNULL)\n    else:\n        try:\n            import py7zr\n        except ImportError:\n            subprocess.run(['pip', 'install', '-q', 'py7zr'], check=True)\n            import py7zr\n        with py7zr.SevenZipFile(archive, 'r') as z:\n            z.extractall(WORK)\n\ndef load_png(path):\n    return np.asarray(Image.open(path).convert('RGB'), dtype=np.uint8)\n\ndef load_split(name, n):\n    cache = f'{WORK}/{name}.npy'\n    if os.path.exists(cache):\n        return np.load(cache)\n    if not os.path.isdir(f'{WORK}/{name}'):\n        t = time.time(); extract(f'{IN_DIR}/{name}.7z'); print(f'extracted {name} in {time.time()-t:.0f}s')\n    paths = [f'{WORK}/{name}/{i}.png' for i in range(1, n + 1)]\n    with ProcessPoolExecutor(os.cpu_count()) as ex:\n        arr = np.stack(list(ex.map(load_png, paths, chunksize=2000)))\n    np.save(cache, arr)\n    return arr\n\nt0 = time.time()\nlabels = pd.read_csv(f'{IN_DIR}/trainLabels.csv').sort_values('id')\nX_all = load_split('train', 50000)                      # (50000, 32, 32, 3) uint8\ny_all = labels['label'].map(CLASSES.index).values.astype(np.int64)\nprint('train', X_all.shape, np.bincount(y_all), f'{time.time()-t0:.0f}s')\n\n# Hold out 5,000 images (stratified) for honest validation\nVAL_SIZE = 5000\nrng = np.random.RandomState(SEED)\nval_idx = np.concatenate([rng.permutation(np.where(y_all == c)[0])[:VAL_SIZE // 10] for c in range(10)])\ntr_mask = np.ones(len(y_all), bool); tr_mask[val_idx] = False\nX_tr, y_tr, X_val, y_val = X_all[tr_mask], y_all[tr_mask], X_all[val_idx], y_all[val_idx]\nprint('train/val:', len(y_tr), len(y_val))\n\nMEAN = torch.tensor(X_tr.reshape(-1, 3).mean(0) / 255, dtype=torch.float32, device=DEVICE).view(1, 3, 1, 1)\nSTD = torch.tensor(X_tr.reshape(-1, 3).std(0) / 255, dtype=torch.float32, device=DEVICE).view(1, 3, 1, 1)\nprint('mean', MEAN.flatten().tolist(), 'std', STD.flatten().tolist())\n\n# %% [Cell 2]\n# ===== 2. Custom network: SE-ResMixNet (designed for 32x32 inputs) =====\n# - two-conv stem keeps full 32x32 resolution\n# - pre-activation residual blocks + Squeeze-and-Excitation channel attention\n# - anti-aliased downsampling shortcut (avg-pool + 1x1 conv)\n# - stochastic depth (DropPath) growing linearly with depth\n# - head: concat(global avg pool, global max pool) -> dropout -> linear\nclass SqueezeExcite(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        w = F.adaptive_avg_pool2d(x, 1)\n        return x * torch.sigmoid(self.fc2(F.silu(self.fc1(w))))\n\nclass DropPath(nn.Module):\n    def __init__(self, p=0.):\n        super().__init__(); self.p = p\n    def forward(self, x):\n        if not self.training or self.p == 0:\n            return x\n        keep = 1 - self.p\n        return x * x.new_empty(x.size(0), 1, 1, 1).bernoulli_(keep) / keep\n\nclass PreActSEBlock(nn.Module):\n    def __init__(self, cin, cout, stride=1, drop=0.):\n        super().__init__()\n        self.bn1 = nn.BatchNorm2d(cin)\n        self.conv1 = nn.Conv2d(cin, cout, 3, stride, 1, bias=False)\n        self.bn2 = nn.BatchNorm2d(cout)\n        self.conv2 = nn.Conv2d(cout, cout, 3, 1, 1, bias=False)\n        self.se = SqueezeExcite(cout)\n        self.drop = DropPath(drop)\n        self.shortcut = None\n        if stride != 1 or cin != cout:\n            self.shortcut = nn.Sequential(nn.AvgPool2d(stride) if stride > 1 else nn.Identity(),\n                                          nn.Conv2d(cin, cout, 1, bias=False))\n    def forward(self, x):\n        out = F.relu(self.bn1(x))\n        sc = x if self.shortcut is None else self.shortcut(out)\n        out = self.conv1(out)\n        out = self.conv2(F.relu(self.bn2(out)))\n        return sc + self.drop(self.se(out))\n\nclass SEResMixNet(nn.Module):\n    def __init__(self, widths=(96, 192, 384), depths=(2, 3, 3), num_classes=10, drop_path=0.1, dropout=0.2):\n        super().__init__()\n        self.stem = nn.Sequential(\n            nn.Conv2d(3, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64), nn.ReLU(inplace=True),\n            nn.Conv2d(64, widths[0], 3, 1, 1, bias=False))\n        n_blocks, k, blocks, cin = sum(depths), 0, [], widths[0]\n        for s, (w, d) in enumerate(zip(widths, depths)):\n            for i in range(d):\n                stride = 2 if (i == 0 and s > 0) else 1\n                blocks.append(PreActSEBlock(cin, w, stride, drop_path * k / (n_blocks - 1)))\n                cin, k = w, k + 1\n        self.blocks = nn.Sequential(*blocks)\n        self.bn = nn.BatchNorm2d(cin)\n        self.dropout = nn.Dropout(dropout)\n        self.fc = nn.Linear(2 * cin, num_classes)\n        for m in self.modules():                      # random initialisation (no pre-trained weights)\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            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.normal_(m.weight, 0, 0.01); nn.init.zeros_(m.bias)\n    def forward(self, x):\n        x = self.blocks(self.stem(x))\n        x = F.relu(self.bn(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.dropout(x))\n\n_m = SEResMixNet()\nprint(f'SEResMixNet parameters: {sum(p.numel() for p in _m.parameters())/1e6:.2f}M')\nprint(_m(torch.randn(2, 3, 32, 32)).shape)\ndel _m\n\n# %% [Cell 3]\n# ===== 3. Training: AdamW + linear warmup -> cosine annealing + CutMix =====\nEPOCHS        = 200\nBATCH         = 512\nPEAK_LR       = 2e-3\nMIN_LR        = 1e-6\nWEIGHT_DECAY  = 0.05\nWARMUP_EPOCHS = 5\nCUTMIX_PROB   = 0.5      # probability that a batch is CutMix-ed\nCUTMIX_ALPHA  = 1.0\nNO_MIX_LAST   = 10       # final epochs without CutMix (fine-tune on clean images)\nLABEL_SMOOTH  = 0.1\nEMA_DECAY     = 0.999\n\n# ---- GPU-side augmentation (fast: whole train set lives on the GPU as uint8) ----\nXtr_gpu = torch.from_numpy(X_tr).permute(0, 3, 1, 2).contiguous().to(DEVICE)\nytr_gpu = torch.from_numpy(y_tr).to(DEVICE)\n\ndef normalize(xb_uint8):\n    return ((xb_uint8.float() / 255.) - MEAN) / STD\n\ndef 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    out = xp[torch.arange(B, device=x.device)[:, None, None, None],\n             torch.arange(C, device=x.device)[None, :, None, None], 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 cutmix(x, y, alpha=CUTMIX_ALPHA):\n    lam = np.random.beta(alpha, alpha)\n    perm = torch.randperm(x.size(0), device=x.device)\n    H, W = x.shape[2:]\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 = np.clip(cy - rh // 2, 0, H), np.clip(cy + rh // 2, 0, H)\n    x1, x2 = np.clip(cx - rw // 2, 0, W), np.clip(cx + rw // 2, 0, 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\n\n# ---- model, EMA copy, optimizer, schedule ----\nmodel = SEResMixNet().to(DEVICE).to(memory_format=torch.channels_last)\nema = copy.deepcopy(model).eval()\nfor p in ema.parameters(): p.requires_grad_(False)\nnet = nn.DataParallel(model) if N_GPU > 1 else model\n\ndecay    = [p for n, p in model.named_parameters() if p.ndim > 1]\nno_decay = [p for n, p in model.named_parameters() if p.ndim <= 1]     # BN & biases: no weight decay\nopt = torch.optim.AdamW([{'params': decay, 'weight_decay': WEIGHT_DECAY},\n                         {'params': no_decay, 'weight_decay': 0.0}], lr=PEAK_LR, betas=(0.9, 0.999))\n\nsteps_per_epoch = len(y_tr) // BATCH\ntotal_steps, warmup_steps = EPOCHS * steps_per_epoch, WARMUP_EPOCHS * steps_per_epoch\ndef lr_factor(step):\n    if step < warmup_steps:                                     # linear warmup\n        return (step + 1) / warmup_steps\n    t = (step - warmup_steps) / max(1, total_steps - warmup_steps)  # cosine annealing\n    return (MIN_LR + (PEAK_LR - MIN_LR) * 0.5 * (1 + math.cos(math.pi * t))) / PEAK_LR\nsched = torch.optim.lr_scheduler.LambdaLR(opt, lr_factor)\nscaler = torch.amp.GradScaler('cuda')\ncriterion = nn.CrossEntropyLoss(label_smoothing=LABEL_SMOOTH)\n\nema_float = [t for t in ema.state_dict().values() if t.dtype.is_floating_point]\nmdl_float = [t for t in model.state_dict().values() if t.dtype.is_floating_point]\nema_other = [t for t in ema.state_dict().values() if not t.dtype.is_floating_point]\nmdl_other = [t for t in model.state_dict().values() if not t.dtype.is_floating_point]\n@torch.no_grad()\ndef update_ema(step):\n    d = min(EMA_DECAY, (1 + step) / (10 + step))\n    torch._foreach_mul_(ema_float, d)\n    torch._foreach_add_(ema_float, mdl_float, alpha=1 - d)\n    for e, m in zip(ema_other, mdl_other): e.copy_(m)\n\n@torch.no_grad()\ndef predict(m, X_uint8_nhwc, bs=2000, tta=True):\n    \"\"\"Softmax probabilities, averaged over original + horizontal flip (TTA).\"\"\"\n    m.eval(); out = []\n    for k in range(0, len(X_uint8_nhwc), bs):\n        xb = torch.from_numpy(X_uint8_nhwc[k:k + bs]).to(DEVICE).permute(0, 3, 1, 2)\n        xb = normalize(xb).contiguous(memory_format=torch.channels_last)\n        with torch.autocast('cuda', dtype=torch.float16):\n            p = m(xb).float().softmax(1)\n            if tta: p = (p + m(xb.flip(3)).float().softmax(1)) / 2\n        out.append(p.cpu())\n    return torch.cat(out).numpy()\n\nhistory, step, t_start = [], 0, time.time()\nfor epoch in range(EPOCHS):\n    net.train(); t_ep = time.time()\n    perm = torch.randperm(len(y_tr), device=DEVICE)\n    use_mix = epoch < EPOCHS - NO_MIX_LAST\n    run_loss, run_correct, seen = 0., 0., 0\n    for b in range(steps_per_epoch):\n        idx = perm[b * BATCH:(b + 1) * BATCH]\n        xb = normalize(random_crop_flip(Xtr_gpu[idx])).contiguous(memory_format=torch.channels_last)\n        yb = ytr_gpu[idx]\n        mixed = use_mix and np.random.rand() < CUTMIX_PROB\n        if mixed:\n            xb, ya, yb2, lam = cutmix(xb, yb)\n        with torch.autocast('cuda', dtype=torch.float16):\n            logits = net(xb)\n            loss = lam * criterion(logits, ya) + (1 - lam) * criterion(logits, yb2) if mixed else criterion(logits, yb)\n        opt.zero_grad(set_to_none=True)\n        scaler.scale(loss).backward()\n        scaler.unscale_(opt)\n        nn.utils.clip_grad_norm_(model.parameters(), 5.0)\n        scaler.step(opt); scaler.update(); sched.step()\n        update_ema(step); step += 1\n        run_loss += loss.item() * len(yb); seen += len(yb)\n        if not mixed:\n            run_correct += (logits.argmax(1) == yb).sum().item()\n    rec = {'epoch': epoch + 1, 'lr': opt.param_groups[0]['lr'], 'train_loss': run_loss / seen,\n           'epoch_time_s': time.time() - t_ep}\n    if (epoch + 1) % 5 == 0 or epoch >= EPOCHS - 10 or epoch == 0:\n        rec['val_acc'] = (predict(model, X_val, tta=False).argmax(1) == y_val).mean()\n        rec['val_acc_ema'] = (predict(ema, X_val, tta=False).argmax(1) == y_val).mean()\n    history.append(rec)\n    print(' | '.join(f'{k}={v:.4f}' if isinstance(v, float) else f'{k}={v}' for k, v in rec.items()),\n          f'| elapsed={(time.time()-t_start)/60:.1f}min', flush=True)\n\nhist = pd.DataFrame(history); hist.to_csv('training_history.csv', index=False)\ntorch.save(model.state_dict(), 'seresmixnet_last.pt'); torch.save(ema.state_dict(), 'seresmixnet_ema.pt')\np_raw, p_ema = predict(model, X_val), predict(ema, X_val)\nacc_raw, acc_ema = (p_raw.argmax(1) == y_val).mean(), (p_ema.argmax(1) == y_val).mean()\nprint(f'Final validation accuracy with flip-TTA: model={acc_raw:.4f}  EMA={acc_ema:.4f}')\nfinal_model = ema if acc_ema >= acc_raw else model\nprint('Using', 'EMA' if final_model is ema else 'raw', 'weights for the submission')\n\n# %% [Cell 4]\n# ===== 4. Learning curves + Kaggle submission =====\nimport matplotlib.pyplot as plt\nfig, ax = plt.subplots(1, 3, figsize=(16, 4))\nax[0].plot(hist.epoch, hist.lr); ax[0].set_title('Learning rate (warmup + cosine)'); ax[0].set_xlabel('epoch')\nax[1].plot(hist.epoch, hist.train_loss); ax[1].set_title('Training loss'); ax[1].set_xlabel('epoch')\nv = hist.dropna(subset=['val_acc'])\nax[2].plot(v.epoch, v.val_acc, 'o-', label='model'); ax[2].plot(v.epoch, v.val_acc_ema, 's-', label='EMA')\nax[2].set_title('Validation accuracy (5k held-out)'); ax[2].set_xlabel('epoch'); ax[2].legend(); ax[2].grid(alpha=.3)\nplt.tight_layout(); plt.savefig('training_curves.png', dpi=120); plt.show()\n\ndel Xtr_gpu; torch.cuda.empty_cache()\nt0 = time.time()\nX_test = load_split('test', 300000)                       # (300000, 32, 32, 3)\nprint('test', X_test.shape, f'loaded in {time.time()-t0:.0f}s')\nprobs = predict(final_model, X_test)\nsub = pd.DataFrame({'id': np.arange(1, len(X_test) + 1), 'label': [CLASSES[k] for k in probs.argmax(1)]})\nsub.to_csv('submission.csv', index=False)\nprint(sub.head(), '\\n', sub.label.value_counts())","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-10-10T11:55:16.289255Z","iopub.execute_input":"2026-10-10T11:55:16.289716Z","iopub.status.idle":"2026-10-10T14:41:59.616465Z","shell.execute_reply.started":"2026-10-10T11:55:16.289687Z","shell.execute_reply":"2026-10-10T14:41:59.615646Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Drawing the Network","metadata":{}},{"cell_type":"code","source":"# Show the SE-ResMixNet architecture\nprint(model)  # full layer-by-layer structure\n\n# Parameter count per part of the network\nfor name, part in [('stem', model.stem), ('blocks', model.blocks),\n                   ('head', nn.ModuleList([model.bn, model.fc]))]:\n    print(f'{name:7s}: {sum(p.numel() for p in part.parameters())/1e6:.2f}M params')\nprint(f'total  : {sum(p.numel() for p in model.parameters())/1e6:.2f}M params')\n\n# Table of output shapes and parameters per layer (installs torchinfo if missing)\ntry:\n    from torchinfo import summary\nexcept ImportError:\n    import subprocess; subprocess.run(['pip', 'install', '-q', 'torchinfo']); from torchinfo import summary\nsummary(model, input_size=(1, 3, 32, 32), depth=3,\n        col_names=('output_size', 'num_params'), device=DEVICE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-10T14:41:59.62063Z","iopub.execute_input":"2026-10-10T14:41:59.621021Z","iopub.status.idle":"2026-10-10T14:42:00.005673Z","shell.execute_reply.started":"2026-10-10T14:41:59.620984Z","shell.execute_reply":"2026-10-10T14:42:00.005096Z"}},"outputs":[],"execution_count":null}]}