{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":3649,"databundleVersionId":46718,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q py7zr\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-05T17:49:13.050683Z","iopub.execute_input":"2025-11-05T17:49:13.051308Z","iopub.status.idle":"2025-11-05T17:49:18.665668Z","shell.execute_reply.started":"2025-11-05T17:49:13.051279Z","shell.execute_reply":"2025-11-05T17:49:18.664929Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# Cell 1: Ultra-fast extraction + config + datasets + loaders\n# ============================================================\n\nimport os, math, random, time, gc\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport py7zr\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torchvision import transforms\n\n# ----------------------------\n# Reproducibility\n# ----------------------------\nSEED = 1337\ndef set_seed(seed=SEED):\n    import random, numpy as np, torch\n    random.seed(seed); np.random.seed(seed)\n    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.benchmark = True\nset_seed(SEED)\n\nprint(\"=\"*70)\nprint(\"ULTRA-FAST EXTRACTION WITH PY7ZR\")\nprint(\"=\"*70)\n\n# ----------------------------\n# Extract train.7z -> /kaggle/working/\n# ----------------------------\ntrain_7z = next((p for p in [\"../input/cifar-10/train.7z\",\n                             \"/kaggle/input/cifar-10/train.7z\"]\n                 if Path(p).exists()), None)\nassert train_7z, \"train.7z not found\"\nwith py7zr.SevenZipFile(train_7z, \"r\") as z:\n    z.extractall(path=\"/kaggle/working/\")\ntrain_dir = Path(\"/kaggle/working/train\")\ntrain_count = len(list(train_dir.glob(\"*.png\")))\nprint(f\"✓ Extracted {train_count} training images -> {train_dir}\")\n\n# ----------------------------\n# Labels\n# ----------------------------\nlabels_csv = next((p for p in [\"../input/cifar-10/trainLabels.csv\",\n                               \"/kaggle/input/cifar-10/trainLabels.csv\"]\n                   if Path(p).exists()), None)\nassert labels_csv, \"trainLabels.csv not found\"\ntrain_labels = pd.read_csv(labels_csv)\nclasses = sorted(train_labels[\"label\"].unique())\nprint(f\"✓ Classes ({len(classes)}): {classes}\")\nprint(f\"✓ Training samples: {len(train_labels):,}\")\n\n# ----------------------------\n# Extract test.7z -> /kaggle/working/\n# ----------------------------\ntest_7z = next((p for p in [\"../input/cifar-10/test.7z\",\n                            \"/kaggle/input/cifar-10/test.7z\"]\n                if Path(p).exists()), None)\nassert test_7z, \"test.7z not found\"\nwith py7zr.SevenZipFile(test_7z, \"r\") as z:\n    z.extractall(path=\"/kaggle/working/\")\n\n# possible layouts\ntest_dir1 = Path(\"/kaggle/working/test\")\ntest_dir2 = Path(\"/kaggle/working/test/test\")\nif test_dir1.exists() and any(test_dir1.glob(\"*.png\")):\n    test_path = test_dir1\nelif test_dir2.exists() and any(test_dir2.glob(\"*.png\")):\n    test_path = test_dir2\nelse:\n    # fallback: find the folder with most PNGs under /kaggle/working\n    pngs = list(Path(\"/kaggle/working\").rglob(\"*.png\"))\n    from collections import Counter\n    parent_counts = Counter(p.parent for p in pngs)\n    test_path, _ = parent_counts.most_common(1)[0]\ntest_count = len(list(test_path.glob(\"*.png\")))\nprint(f\"✓ Extracted {test_count} test images -> {test_path}\")\nprint(\"✓ Examples exist:\", (test_path / \"1.png\").exists(), (test_path / \"300000.png\").exists())\n\nprint(\"\\n\" + \"=\"*70)\nprint(\"✅ EXTRACTION COMPLETE\")\nprint(\"=\"*70)\n\n# ----------------------------\n# Config\n# ----------------------------\nclass CFG:\n    img_size = 32\n    num_classes = 10\n    epochs = 120          # bump to 100-120 for ≥0.94\n    batch_size = 512\n    lr = 3e-3\n    weight_decay = 1e-4\n    warmup_epochs = 5\n    num_workers = 4\n    use_amp = True\n    cutmix_alpha = 1.0\n    cutmix_prob  = 0.5\n    cutmix_warmup_epochs = 5   # CutMix off for first N epochs\n    val_ratio = 0.05\n    ema_decay = 0.999\n    tta_times = 4\n    min_lr = 1e-5\n\nCFG.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", CFG.device)\n\n# ----------------------------\n# Class maps\n# ----------------------------\ncls2idx = {c:i for i,c in enumerate(classes)}\nidx2cls = {i:c for c,i in cls2idx.items()}\n\n# ----------------------------\n# Datasets\n# ----------------------------\nclass Cifar10Train(Dataset):\n    def __init__(self, df, img_dir, transform=None):\n        self.df = df.reset_index(drop=True)\n        self.dir = Path(img_dir)\n        self.transform = transform\n    def __len__(self): return len(self.df)\n    def __getitem__(self, i):\n        r = self.df.iloc[i]\n        img = Image.open(self.dir / f\"{r['id']}.png\").convert(\"RGB\")\n        if self.transform: img = self.transform(img)\n        return img, cls2idx[r[\"label\"]]\n\nclass Cifar10Test(Dataset):\n    def __init__(self, img_dir, transform=None):\n        self.dir = Path(img_dir)\n        self.ids = sorted([p.stem for p in self.dir.glob(\"*.png\")], key=lambda x:int(x))\n        self.transform = transform\n    def __len__(self): return len(self.ids)\n    def __getitem__(self, i):\n        img_id = self.ids[i]\n        img = Image.open(self.dir / f\"{img_id}.png\").convert(\"RGB\")\n        if self.transform: img = self.transform(img)\n        return img, img_id\n\n# ----------------------------\n# Transforms  (IMPORTANT: normalization added)\n# ----------------------------\nCIFAR_MEAN = [0.4914, 0.4822, 0.4465]\nCIFAR_STD  = [0.2470, 0.2435, 0.2616]\n\ntrain_tfms = transforms.Compose([\n    transforms.RandomHorizontalFlip(),\n    transforms.RandomCrop(CFG.img_size, padding=4, padding_mode=\"reflect\"),\n    transforms.AutoAugment(transforms.AutoAugmentPolicy.CIFAR10),\n    transforms.ToTensor(),\n    transforms.Normalize(CIFAR_MEAN, CIFAR_STD),\n])\n\ntest_tfms = transforms.Compose([\n    transforms.ToTensor(),\n    transforms.Normalize(CIFAR_MEAN, CIFAR_STD),\n])\n\n# ----------------------------\n# CutMix helpers + collate\n# ----------------------------\ndef rand_bbox(W, H, lam):\n    import math, numpy as np\n    cut_rat = math.sqrt(1. - lam)\n    cut_w = int(W * cut_rat); cut_h = int(H * cut_rat)\n    cx = np.random.randint(W); cy = np.random.randint(H)\n    x1 = np.clip(cx - cut_w//2, 0, W); y1 = np.clip(cy - cut_h//2, 0, H)\n    x2 = np.clip(cx + cut_w//2, 0, W); y2 = np.clip(cy + cut_h//2, 0, H)\n    return x1, y1, x2, y2\n\ndef cutmix_collate(batch):\n    import numpy as np, torch\n    images = torch.stack([b[0] for b in batch], dim=0)\n    targets = torch.tensor([b[1] for b in batch], dtype=torch.long)\n    if np.random.rand() < CFG.cutmix_prob and CFG.cutmix_alpha > 0:\n        lam = np.random.beta(CFG.cutmix_alpha, CFG.cutmix_alpha)\n        B, C, H, W = images.size()\n        idx = torch.randperm(B)\n        x1,y1,x2,y2 = rand_bbox(W, H, lam)\n        images[:, :, y1:y2, x1:x2] = images[idx, :, y1:y2, x1:x2]\n        lam = 1 - ((x2-x1)*(y2-y1)/(W*H))\n        return images, (targets, targets[idx], lam)\n    else:\n        return images, (targets, targets, 1.0)\n\n# ----------------------------\n# Split + loaders\n# ----------------------------\nfull_train_df = train_labels.copy()\nn_train = int(len(full_train_df) * (1 - CFG.val_ratio))\nn_val   = len(full_train_df) - n_train\n\ntr_idx, va_idx = random_split(\n    range(len(full_train_df)), [n_train, n_val],\n    generator=torch.Generator().manual_seed(SEED)\n)\ntrain_df = full_train_df.iloc[tr_idx.indices].reset_index(drop=True)\nval_df   = full_train_df.iloc[va_idx.indices].reset_index(drop=True)\n\nTRAIN_DIR = train_dir\nTEST_DIR  = test_path\n\ntrain_ds = Cifar10Train(train_df, TRAIN_DIR, transform=train_tfms)\nval_ds   = Cifar10Train(val_df,   TRAIN_DIR, transform=test_tfms)\n\ntrain_loader = DataLoader(train_ds, batch_size=CFG.batch_size, shuffle=True,\n                          num_workers=CFG.num_workers, pin_memory=True,\n                          collate_fn=cutmix_collate, drop_last=True)\nval_loader   = DataLoader(val_ds, batch_size=CFG.batch_size, shuffle=False,\n                          num_workers=CFG.num_workers, pin_memory=True)\n\nprint(f\"Train/Val sizes: {len(train_ds)} / {len(val_ds)}\")\nprint(\"👉 Run Cell 2 to train, then Cell 3 to create submission.csv\")\n\n# (Optional) tiny overfit sanity loader (uncomment in Cell 2 to use)\ntiny_idx = list(range(min(256, len(train_ds))))\ntiny_loader = DataLoader(torch.utils.data.Subset(train_ds, tiny_idx),\n                         batch_size=128, shuffle=True, num_workers=2)\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-05T17:49:56.448250Z","iopub.execute_input":"2025-11-05T17:49:56.448823Z","iopub.status.idle":"2025-11-05T17:52:24.769499Z","shell.execute_reply.started":"2025-11-05T17:49:56.448793Z","shell.execute_reply":"2025-11-05T17:52:24.768420Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"## check the test file only \nfrom pathlib import Path\nimport re\n\ndef find_test_dir():\n    cand = [Path(\"/kaggle/working/test\"), Path(\"/kaggle/working/test/test\")]\n    for p in cand:\n        if p.exists() and any(p.glob(\"*.png\")):\n            # Make sure files are numeric 1..300000\n            sample = next((x for x in p.glob(\"*.png\")), None)\n            if sample and re.fullmatch(r\"\\d+\\.png\", sample.name):\n                return p\n    # fallback: search the tree\n    all_png_dirs = {}\n    for png in Path(\"/kaggle/working\").rglob(\"*.png\"):\n        if re.fullmatch(r\"\\d+\\.png\", png.name):\n            all_png_dirs[png.parent] = all_png_dirs.get(png.parent, 0) + 1\n    assert all_png_dirs, \"No numeric test PNGs found after extraction.\"\n    # choose the dir with most numeric pngs\n    test_dir = max(all_png_dirs, key=all_png_dirs.get)\n    return test_dir\n\nTEST_DIR = find_test_dir()\nprint(\"TEST_DIR =\", TEST_DIR)\nprint(\"Has 1.png? \", (TEST_DIR / \"1.png\").exists(), \" | Has 300000.png? \", (TEST_DIR / \"300000.png\").exists())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-05T17:52:30.122870Z","iopub.execute_input":"2025-11-05T17:52:30.123574Z","iopub.status.idle":"2025-11-05T17:52:30.551843Z","shell.execute_reply.started":"2025-11-05T17:52:30.123534Z","shell.execute_reply":"2025-11-05T17:52:30.551253Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# Cell 2: Custom model + training (AdamW, cosine+warmup, CutMix, AMP, EMA)\n# Depends on Cell 1 globals: CFG, train_loader, val_loader, tiny_loader (optional)\n# ============================================================\n\nimport math, time\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\n\n# ------------------------------------------------------------\n# Toggle: evaluate with EMA or live model (start with live!)\n# ------------------------------------------------------------\nEVAL_WITH_EMA = False   # set True later (e.g., after epoch 20) if you want\n\n# ------------------------------------------------------------\n# Model\n# ------------------------------------------------------------\nclass 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        w = F.adaptive_avg_pool2d(x, 1)\n        w = F.relu(self.fc1(w), inplace=True)\n        w = torch.sigmoid(self.fc2(w))\n        return x * w\n\nclass XBlock(nn.Module):\n    def __init__(self, in_c, out_c, stride=1, expansion=2):\n        super().__init__()\n        mid = out_c * expansion\n        self.use_res = (in_c == out_c and stride == 1)\n        self.pw1 = nn.Conv2d(in_c, mid, 1, bias=False); self.bn1 = nn.BatchNorm2d(mid)\n        self.dw  = nn.Conv2d(mid, mid, 3, stride=stride, padding=1, groups=mid, bias=False); self.bn2 = nn.BatchNorm2d(mid)\n        self.pw2 = nn.Conv2d(mid, out_c, 1, bias=False); self.bn3 = nn.BatchNorm2d(out_c)\n        self.se  = SE(out_c, r=8); self.act = nn.ReLU(inplace=True)\n        self.shortcut = nn.Identity() if self.use_res else nn.Sequential(\n            nn.Conv2d(in_c, out_c, 1, stride=stride, bias=False), nn.BatchNorm2d(out_c)\n        )\n    def forward(self, x):\n        out = self.act(self.bn1(self.pw1(x)))\n        out = self.act(self.bn2(self.dw(out)))\n        out = self.bn3(self.pw2(out))\n        out = self.se(out)\n        return F.relu(out + self.shortcut(x), inplace=True)\n\nclass CifarXNet(nn.Module):\n    def __init__(self, num_classes=10, width=64):\n        super().__init__()\n        self.stem = nn.Sequential(\n            nn.Conv2d(3, width, 3, padding=1, bias=False),\n            nn.BatchNorm2d(width),\n            nn.ReLU(inplace=True),\n        )\n        self.stage1 = nn.Sequential(XBlock(width, width), XBlock(width, width))\n        self.stage2 = nn.Sequential(XBlock(width, width*2, stride=2), XBlock(width*2, width*2))\n        self.stage3 = nn.Sequential(XBlock(width*2, width*4, stride=2), XBlock(width*4, width*4))\n        self.head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(width*4, num_classes))\n        self._init_weights()\n    def _init_weights(self):\n        # REQUIREMENT: Randomly initialized weights (no pretraining)\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.Linear):\n                nn.init.kaiming_normal_(m.weight, 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    def forward(self, x):\n        x = self.stem(x); x = self.stage1(x); x = self.stage2(x); x = self.stage3(x)\n        return self.head(x)\n\n# ------------------------------------------------------------\n# EMA\n# ------------------------------------------------------------\nclass ModelEMA:\n    def __init__(self, model, decay=0.999):\n        self.ema = CifarXNet(num_classes=CFG.num_classes).to(CFG.device)\n        self.ema.load_state_dict(model.state_dict())\n        self.decay = decay\n        for p in self.ema.parameters(): p.requires_grad_(False)\n    @torch.no_grad()\n    def update(self, model):\n        d = self.decay; msd = model.state_dict()\n        for k, v in self.ema.state_dict().items():\n            if v.dtype.is_floating_point: v.copy_(v * d + (1. - d) * msd[k])\n\ndef build_warmup_cosine_scheduler(optim, num_epochs, warmup_epochs, min_lr, base_lr):\n    def lr_lambda(e):\n        if e < warmup_epochs: return (e + 1) / max(1, warmup_epochs)\n        p = (e - warmup_epochs) / max(1, num_epochs - warmup_epochs)\n        cos = 0.5 * (1 + math.cos(math.pi * p))\n        return (min_lr / base_lr) + (1 - (min_lr / base_lr)) * cos\n    return torch.optim.lr_scheduler.LambdaLR(optim, lr_lambda=lr_lambda)\n\ncriterion = nn.CrossEntropyLoss()\n\ndef cutmix_criterion(pred, targets_tuple):\n    y1, y2, lam = targets_tuple\n    return lam * criterion(pred, y1) + (1 - lam) * criterion(pred, y2)\n\ndef accuracy(output, target):\n    return (output.argmax(1) == target).float().mean().item()\n\n# AMP compatibility (old/new PyTorch)\nGradScaler = torch.cuda.amp.GradScaler if hasattr(torch.cuda, \"amp\") else torch.amp.GradScaler\namp_context = torch.cuda.amp.autocast if hasattr(torch.cuda.amp, \"autocast\") else torch.amp.autocast\n\ndef evaluate(model, ema, loader):\n    net = ema.ema if EVAL_WITH_EMA else model\n    net.eval()\n    losses, accs = [], []\n    with torch.no_grad(), amp_context(enabled=CFG.use_amp):\n        for batch in loader:\n            # val loader yields (images, targets)\n            if len(batch) == 2 and not isinstance(batch[1], (tuple, list)):\n                images, targets = batch\n            else:\n                images, targets = batch[0], batch[1][0]\n            images, targets = images.to(CFG.device), targets.to(CFG.device)\n            logits = net(images)\n            losses.append(criterion(logits, targets).item())\n            accs.append(accuracy(logits, targets))\n    return float(np.mean(losses)), float(np.mean(accs))\n\n# ------------------------------------------------------------\n# Train\n# ------------------------------------------------------------\nmodel = CifarXNet(num_classes=CFG.num_classes).to(CFG.device)\noptimizer = torch.optim.AdamW(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay)\nscheduler = build_warmup_cosine_scheduler(optimizer, CFG.epochs, CFG.warmup_epochs, CFG.min_lr, CFG.lr)\nscaler = GradScaler(enabled=CFG.use_amp)\nema = ModelEMA(model, decay=CFG.ema_decay)\n\nbest_acc = 0.0\nfor epoch in range(CFG.epochs):\n    model.train()\n    total_loss = total_acc = n = 0\n    start = time.time()\n\n    # CutMix warm-up: disable for first N epochs\n    use_cutmix = (epoch >= CFG.cutmix_warmup_epochs)\n\n    for images, targets_tuple in train_loader:\n        images = images.to(CFG.device)\n        if not use_cutmix:\n            y = targets_tuple[0]\n            targets_tuple = (y, y, 1.0)\n        targets_tuple = tuple(t.to(CFG.device) if torch.is_tensor(t) else t for t in targets_tuple)\n\n        optimizer.zero_grad(set_to_none=True)\n        with amp_context(enabled=CFG.use_amp):\n            logits = model(images)\n            loss = cutmix_criterion(logits, targets_tuple)\n        scaler.scale(loss).backward()\n        scaler.step(optimizer); scaler.update()\n        ema.update(model)\n\n        total_loss += loss.item()\n        total_acc  += accuracy(logits.detach(), targets_tuple[0]); n += 1\n\n    scheduler.step()\n    val_loss, val_acc = evaluate(model, ema, val_loader)\n    lr_now = scheduler.get_last_lr()[0]\n    print(f\"Epoch {epoch+1:03d}/{CFG.epochs} | lr={lr_now:.6f} | \"\n          f\"train_loss={total_loss/n:.4f} train_acc={total_acc/n:.4f} | \"\n          f\"val_loss={val_loss:.4f} val_acc={val_acc:.4f} | {time.time()-start:.1f}s\")\n\n    if val_acc > best_acc:\n        best_acc = val_acc\n        torch.save(model.state_dict(), \"best.pth\")        # live model\n        torch.save(ema.ema.state_dict(), \"best_ema.pth\")  # optional\n\nprint(f\"\\n✅ Best val_acc (live model): {best_acc:.4f}\")\nprint(\"💾 Saved checkpoints: best.pth (live), best_ema.pth (EMA)\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-05T17:55:05.742496Z","iopub.execute_input":"2025-11-05T17:55:05.743220Z","iopub.status.idle":"2025-11-05T18:21:33.887750Z","shell.execute_reply.started":"2025-11-05T17:55:05.743197Z","shell.execute_reply":"2025-11-05T18:21:33.886894Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# Cell 3: Load best (live) model, TTA inference, write submission.csv (300,000 rows)\n# Depends on Cell 1 globals: CFG, Cifar10Test, TEST_DIR, test_tfms, idx2cls\n# ============================================================\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport pandas as pd\n\n# --- redeclare model (safe if already defined) ---\nclass SE(nn.Module):\n    def __init__(self, c, r=8):\n        super().__init__()\n        self.fc1 = nn.Conv2d(c, c // r, 1); self.fc2 = nn.Conv2d(c // r, c, 1)\n    def forward(self, x):\n        w = F.adaptive_avg_pool2d(x, 1); w = F.relu(self.fc1(w), inplace=True)\n        w = torch.sigmoid(self.fc2(w)); return x * w\n\nclass XBlock(nn.Module):\n    def __init__(self, in_c, out_c, stride=1, expansion=2):\n        super().__init__()\n        mid = out_c * expansion\n        self.use_res = (in_c == out_c and stride == 1)\n        self.pw1 = nn.Conv2d(in_c, mid, 1, bias=False); self.bn1 = nn.BatchNorm2d(mid)\n        self.dw  = nn.Conv2d(mid, mid, 3, stride=stride, padding=1, groups=mid, bias=False); self.bn2 = nn.BatchNorm2d(mid)\n        self.pw2 = nn.Conv2d(mid, out_c, 1, bias=False); self.bn3 = nn.BatchNorm2d(out_c)\n        self.se  = SE(out_c, r=8); self.act = nn.ReLU(inplace=True)\n        self.shortcut = nn.Identity() if self.use_res else nn.Sequential(\n            nn.Conv2d(in_c, out_c, 1, stride=stride, bias=False), nn.BatchNorm2d(out_c)\n        )\n    def forward(self, x):\n        out = self.act(self.bn1(self.pw1(x)))\n        out = self.act(self.bn2(self.dw(out)))\n        out = self.bn3(self.pw2(out)); out = self.se(out)\n        return F.relu(out + self.shortcut(x), inplace=True)\n\nclass CifarXNet(nn.Module):\n    def __init__(self, num_classes=10, width=64):\n        super().__init__()\n        self.stem = nn.Sequential(\n            nn.Conv2d(3, width, 3, padding=1, bias=False),\n            nn.BatchNorm2d(width),\n            nn.ReLU(inplace=True),\n        )\n        self.stage1 = nn.Sequential(XBlock(width, width), XBlock(width, width))\n        self.stage2 = nn.Sequential(XBlock(width, width*2, stride=2), XBlock(width*2, width*2))\n        self.stage3 = nn.Sequential(XBlock(width*2, width*4, stride=2), XBlock(width*4, width*4))\n        self.head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(width*4, CFG.num_classes))\n        self._init_weights()\n    def _init_weights(self):\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode=\"fan_out\", nonlinearity=\"relu\")\n            elif isinstance(m, nn.Linear):\n                nn.init.kaiming_normal_(m.weight, 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    def forward(self, x):\n        x = self.stem(x); x = self.stage1(x); x = self.stage2(x); x = self.stage3(x)\n        return self.head(x)\n\n# AMP compatibility\namp_context = torch.cuda.amp.autocast if hasattr(torch.cuda.amp, \"autocast\") else torch.amp.autocast\n\n# Build test loader\ntest_ds = Cifar10Test(TEST_DIR, transform=test_tfms)\ntest_loader = DataLoader(test_ds, batch_size=CFG.batch_size, shuffle=False,\n                         num_workers=CFG.num_workers, pin_memory=True)\n\n# Load best *live* model checkpoint (not EMA)\nmodel = CifarXNet(num_classes=CFG.num_classes).to(CFG.device)\nstate = torch.load(\"best.pth\", map_location=\"cpu\")\nmodel.load_state_dict(state); model.eval()\n\n@torch.no_grad()\ndef predict_tta(model, loader, tta_times=4):\n    probs_list, ids_all = [], []\n    for images, ids in loader:\n        images = images.to(CFG.device)\n        with amp_context(enabled=CFG.use_amp):\n            logits = model(images)\n            prob = F.softmax(logits, dim=1)\n            for _ in range(tta_times - 1):\n                flipped = torch.flip(images, dims=[3])  # horizontal flip\n                prob += F.softmax(model(flipped), dim=1)\n            prob /= tta_times\n        probs_list.append(prob.cpu()); ids_all.extend(ids)\n    probs = torch.cat(probs_list, dim=0)\n    preds = probs.argmax(1).numpy()\n    return ids_all, preds\n\ntest_ids, test_preds_idx = predict_tta(model, test_loader, tta_times=CFG.tta_times)\ntest_preds = [idx2cls[int(i)] for i in test_preds_idx]\n\nsub = pd.DataFrame({\"id\": test_ids, \"label\": test_preds})\nsub[\"id\"] = sub[\"id\"].astype(int)\nsub = sub.sort_values(\"id\").reset_index(drop=True)\nassert len(sub) == 300_000, f\"Submission must have 300,000 rows, got {len(sub)}\"\n\nsub.to_csv(\"submission.csv\", index=False)\nprint(sub.head(), \"\\n\"); print(sub.tail(), \"\\n\")\nprint(f\"✅ Wrote submission.csv with {len(sub)} rows\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-05T18:23:49.637132Z","iopub.execute_input":"2025-11-05T18:23:49.637880Z","iopub.status.idle":"2025-11-05T18:26:04.677727Z","shell.execute_reply.started":"2025-11-05T18:23:49.637850Z","shell.execute_reply":"2025-11-05T18:26:04.676876Z"}},"outputs":[],"execution_count":null}]}