{"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"},{"sourceId":312126,"sourceType":"datasetVersion","datasetId":130727}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ===========================\n# Cell 1: Imports • Config • Transforms • Extract (.7z) • Datasets • Loaders\n# ===========================\nimport os, random\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader, Subset\nfrom torchvision import transforms\nfrom torchvision.datasets import CIFAR10\n\n# ----- Repro & Device -----\nSEED = 1337\nrandom.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ntorch.backends.cudnn.benchmark = True\n\n# ----- Kaggle roots (robust detection) -----\nCANDIDATES = [\n    Path(\"/kaggle/input/cifar-10\"),\n    Path(\"/kaggle/input/cifar10-object-recognition-in-images-zip-file\"),\n]\nKAGGLE_ROOT = next((p for p in CANDIDATES if p.exists()), CANDIDATES[0])\nUSE_KAGGLE  = KAGGLE_ROOT.exists()\n\nWORK_ROOT   = Path(\"/kaggle/working/cifar10\")  # where we extract\nWORK_TRAIN  = WORK_ROOT / \"train\"\nWORK_TEST   = WORK_ROOT / \"test\"\n\n# ----- Labels -----\nLABELS = [\"airplane\",\"automobile\",\"bird\",\"cat\",\"deer\",\"dog\",\"frog\",\"horse\",\"ship\",\"truck\"]\nlabel2id = {c:i for i,c in enumerate(LABELS)}\nid2label = {i:c for c,i in label2id.items()}\n\n# ----- Transforms -----\ntrain_tfms = transforms.Compose([\n    transforms.RandomCrop(32, padding=4, padding_mode=\"reflect\"),\n    transforms.RandomHorizontalFlip(),\n    transforms.ColorJitter(0.2,0.2,0.2,0.1),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=(0.4914,0.4822,0.4465), std=(0.2023,0.1994,0.2010)),\n    transforms.RandomErasing(p=0.5, scale=(0.02,0.15), ratio=(0.3,3.3), value=\"random\"),\n])\ntest_tfms = transforms.Compose([\n    transforms.ToTensor(),\n    transforms.Normalize(mean=(0.4914,0.4822,0.4465), std=(0.2023,0.1994,0.2010)),\n])\n\nVAL_RATIO   = 0.01  # ← smaller val, more train data\nNUM_WORKERS = max(2, (os.cpu_count() or 4) // 2)\n\ndef make_loader(ds, bs, shuffle):\n    return DataLoader(\n        ds, batch_size=bs, shuffle=shuffle,\n        num_workers=NUM_WORKERS, pin_memory=True,\n        persistent_workers=(NUM_WORKERS > 0), drop_last=shuffle\n    )\n\n# --------------------------\n# Extract helper for .7z archives\n# --------------------------\ndef ensure_extracted():\n    WORK_ROOT.mkdir(parents=True, exist_ok=True)\n    train_pngs = list(WORK_TRAIN.glob(\"*.png\"))\n    test_pngs  = list(WORK_TEST.glob(\"*.png\"))\n\n    need_train = len(train_pngs) == 0\n    need_test  = len(test_pngs)  == 0\n\n    if not (need_train or need_test):\n        print(f\"[extract] Already extracted: train={len(train_pngs)} test={len(test_pngs)}\")\n        # still flatten in case previous run left nested dirs\n    # Try import py7zr if archives exist and we need extraction\n    if (KAGGLE_ROOT / \"train.7z\").exists() or (KAGGLE_ROOT / \"test.7z\").exists():\n        if need_train or need_test:\n            try:\n                import py7zr  # noqa\n            except Exception:\n                import sys, subprocess\n                print(\"[extract] Installing py7zr …\")\n                subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"py7zr\"])\n                import py7zr  # noqa\n\n        if need_train and (KAGGLE_ROOT / \"train.7z\").exists():\n            src = KAGGLE_ROOT / \"train.7z\"\n            WORK_TRAIN.mkdir(parents=True, exist_ok=True)\n            print(\"[extract] Extracting train.7z →\", WORK_TRAIN)\n            import py7zr\n            with py7zr.SevenZipFile(src, mode='r') as z:\n                z.extractall(path=WORK_TRAIN)\n\n        if need_test and (KAGGLE_ROOT / \"test.7z\").exists():\n            src = KAGGLE_ROOT / \"test.7z\"\n            WORK_TEST.mkdir(parents=True, exist_ok=True)\n            print(\"[extract] Extracting test.7z →\", WORK_TEST)\n            import py7zr\n            with py7zr.SevenZipFile(src, mode='r') as z:\n                z.extractall(path=WORK_TEST)\n\n    # Flatten if nested dirs exist (applies even if already extracted)\n    for split_dir in [WORK_TRAIN, WORK_TEST]:\n        if not split_dir.exists():  # nothing to do yet\n            continue\n        subdirs = [p for p in split_dir.iterdir() if p.is_dir()]\n        if subdirs:\n            for sd in subdirs:\n                for p in sd.rglob(\"*.png\"):\n                    target = split_dir / p.name\n                    if not target.exists():\n                        p.replace(target)\n            for sd in subdirs:\n                try: sd.rmdir()\n                except Exception: pass\n\nif USE_KAGGLE:\n    ensure_extracted()\n\n# --------------------------\n# Datasets (competition style: flat train/*.png and test/*.png + CSV)\n# --------------------------\nif USE_KAGGLE:\n    csv_path = KAGGLE_ROOT / \"trainLabels.csv\"\n    assert csv_path.exists(), f\"Missing labels CSV at {csv_path}\"\n\n    df_all = pd.read_csv(csv_path)\n    id2path = {p.stem: p for p in WORK_TRAIN.glob(\"*.png\")}\n    before = len(df_all)\n    df_all = df_all[df_all[\"id\"].astype(str).isin(id2path.keys())].reset_index(drop=True)\n    missing = before - len(df_all)\n    print(f\"[cifar-10] kept {len(df_all)}/{before} | missing files after extract: {missing}\")\n    assert len(df_all) > 0, \"No train images found after extraction.\"\n\n    # split\n    idx = np.arange(len(df_all)); np.random.shuffle(idx)\n    v = int(len(idx) * VAL_RATIO)\n    val_idx, tr_idx = idx[:v], idx[v:]\n\n    class CSVCIFARTrain(Dataset):\n        def __init__(self, df, id2path, transform):\n            self.df, self.id2path, self.transform = df, id2path, transform\n        def __len__(self): return len(self.df)\n        def __getitem__(self, i):\n            row = self.df.iloc[i]\n            pid = str(row[\"id\"])\n            img = Image.open(self.id2path[pid]).convert(\"RGB\")\n            return self.transform(img), label2id[row[\"label\"]]\n\n    train_ds = CSVCIFARTrain(df_all.iloc[tr_idx].reset_index(drop=True), id2path, train_tfms)\n    val_ds   = CSVCIFARTrain(df_all.iloc[val_idx].reset_index(drop=True), id2path, test_tfms)\n\n    test_ids = sorted([p.stem for p in WORK_TEST.glob(\"*.png\")])  # sorted for consistency\n\n    class FolderTest(Dataset):\n        def __init__(self, ids, base, transform):\n            self.ids, self.base, self.transform = ids, base, transform\n        def __len__(self): return len(self.ids)\n        def __getitem__(self, i):\n            pid = self.ids[i]\n            img = Image.open(self.base / f\"{pid}.png\").convert(\"RGB\")\n            return self.transform(img), pid\n\n    test_ds = FolderTest(test_ids, WORK_TEST, test_tfms)\n\nelse:\n    # Local/debug fallback (will require internet if not cached)\n    train_all = CIFAR10(root=\"./\", train=True,  download=True, transform=train_tfms)\n    val_all   = CIFAR10(root=\"./\", train=True,  download=True, transform=test_tfms)\n    n = len(train_all)\n    idx = np.arange(n); np.random.shuffle(idx)\n    v = int(n * VAL_RATIO)\n    val_idx, tr_idx = idx[:v], idx[v:]\n    train_ds, val_ds = Subset(train_all, tr_idx), Subset(val_all, val_idx)\n    test_ds = CIFAR10(root=\"./\", train=False, download=True, transform=test_tfms)\n\n# ----- Build loaders (with auto batch and fallback) -----\nBATCH = 384  # safe for P100 16GB; will fallback if needed\ntry:\n    train_dl = make_loader(train_ds, BATCH, True)\n    val_dl   = make_loader(val_ds,   BATCH, False)\n    test_dl  = make_loader(test_ds,  BATCH, False)\nexcept RuntimeError:\n    BATCH = 256\n    train_dl = make_loader(train_ds, BATCH, True)\n    val_dl   = make_loader(val_ds,   BATCH, False)\n    test_dl  = make_loader(test_ds,  BATCH, False)\n\nprint(\"Device:\", device, \"| Batch:\", BATCH, \"| Using Kaggle:\", USE_KAGGLE, \"| Root:\", str(KAGGLE_ROOT))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-11T15:31:08.773596Z","iopub.execute_input":"2025-10-11T15:31:08.774417Z","iopub.status.idle":"2025-10-11T15:31:14.123621Z","shell.execute_reply.started":"2025-10-11T15:31:08.774388Z","shell.execute_reply":"2025-10-11T15:31:14.122724Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ===== Cell 2 (replace): GhanimNetV2 — Wide PreAct-ResNet + SE + DropPath =====\nimport math\nimport torch\nimport torch.nn as nn\n\ndef conv3x3(c_in, c_out, stride=1, groups=1):\n    return nn.Conv2d(c_in, c_out, kernel_size=3, stride=stride, padding=1,\n                     groups=groups, bias=False)\n\ndef conv1x1(c_in, c_out, stride=1):\n    return nn.Conv2d(c_in, c_out, kernel_size=1, stride=stride, bias=False)\n\n# --- Stochastic Depth (DropPath) ---\nclass DropPath(nn.Module):\n    def __init__(self, p: float = 0.0):\n        super().__init__()\n        self.p = float(p)\n    def forward(self, x):\n        if not self.training or self.p == 0.0:\n            return x\n        keep = 1 - self.p\n        shape = (x.size(0),) + (1,) * (x.ndim - 1)\n        mask = x.new_empty(shape).bernoulli_(keep)\n        return x / keep * mask\n\n# --- Squeeze & Excitation ---\nclass SE(nn.Module):\n    def __init__(self, ch, r=16):\n        super().__init__()\n        self.pool = nn.AdaptiveAvgPool2d(1)\n        self.fc = nn.Sequential(\n            nn.Conv2d(ch, max(8, ch // r), 1, bias=True), nn.ReLU(inplace=True),\n            nn.Conv2d(max(8, ch // r), ch, 1, bias=True), nn.Sigmoid()\n        )\n    def forward(self, x):\n        return x * self.fc(self.pool(x))\n\n# --- Pre-activation Basic Residual Block ---\nclass PreActBlock(nn.Module):\n    def __init__(self, c_in, c_out, stride=1, se=False, drop_path=0.0):\n        super().__init__()\n        self.bn1  = nn.BatchNorm2d(c_in)\n        self.act1 = nn.ReLU(inplace=True)\n        self.conv1 = conv3x3(c_in, c_out, stride=stride)\n\n        self.bn2  = nn.BatchNorm2d(c_out)\n        self.act2 = nn.ReLU(inplace=True)\n        self.conv2 = conv3x3(c_out, c_out, stride=1)\n\n        self.se   = SE(c_out) if se else nn.Identity()\n        self.down = conv1x1(c_in, c_out, stride) if (stride != 1 or c_in != c_out) else None\n        self.drop = DropPath(drop_path)\n\n    def forward(self, x):\n        idt = x\n        out = self.act1(self.bn1(x))\n        if self.down is not None:\n            idt = self.down(out)                   # pre-act shortcut\n        out = self.conv1(out)\n        out = self.conv2(self.act2(self.bn2(out)))\n        out = self.se(out)\n        out = self.drop(out)\n        return out + idt\n\n# --- Model ---\nclass GhanimNetV2(nn.Module):\n    \"\"\"\n    CIFAR-10 custom net:\n      stem -> [stage1, stage2, stage3, stage4] with depths [3,4,6,9] (last stage deeper)\n      widths: [96, 192, 384, 512]\n      PreAct residual + SE + DropPath\n    \"\"\"\n    def __init__(self, num_classes=10, se_every=2, drop_path_max=0.15, widths=(96,192,384,512), depths=(3,4,6,9)):\n        super().__init__()\n\n        # CIFAR-friendly stem (3x3)\n        self.stem = nn.Sequential(\n            nn.Conv2d(3, widths[0], kernel_size=3, stride=1, padding=1, bias=False),\n            nn.BatchNorm2d(widths[0]),\n            nn.ReLU(inplace=True)\n        )\n\n        # total blocks to linearly scale drop_path across depth\n        total_blocks = sum(depths)\n        def make_stage(c_in, c_out, n, stride, start_idx):\n            blocks = []\n            for i in range(n):\n                dp = drop_path_max * (start_idx + i) / max(1, total_blocks - 1)\n                blocks.append(\n                    PreActBlock(\n                        c_in if i == 0 else c_out,\n                        c_out,\n                        stride=stride if i == 0 else 1,\n                        se=((i % se_every) == (se_every - 1)),\n                        drop_path=dp\n                    )\n                )\n            return nn.Sequential(*blocks)\n\n        # build stages\n        idx = 0\n        self.layer1 = make_stage(widths[0], widths[0], depths[0], stride=1, start_idx=idx); idx += depths[0]\n        self.layer2 = make_stage(widths[0], widths[1], depths[1], stride=2, start_idx=idx); idx += depths[1]\n        self.layer3 = make_stage(widths[1], widths[2], depths[2], stride=2, start_idx=idx); idx += depths[2]\n        self.layer4 = make_stage(widths[2], widths[3], depths[3], stride=2, start_idx=idx); idx += depths[3]\n\n        self.bn_head = nn.BatchNorm2d(widths[3])\n        self.act_head = nn.ReLU(inplace=True)\n        self.pool = nn.AdaptiveAvgPool2d(1)\n        self.drop = nn.Dropout(p=0.2)\n        self.fc = nn.Linear(widths[3], num_classes)\n\n        # Kaiming init\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.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\n    def forward(self, x):\n        x = self.stem(x)\n        x = self.layer1(x); x = self.layer2(x); x = self.layer3(x); x = self.layer4(x)\n        x = self.act_head(self.bn_head(x))\n        x = self.pool(x).flatten(1)\n        x = self.drop(x)\n        return self.fc(x)\n\n# instantiate\nmodel = GhanimNetV2(num_classes=10).to(device).to(memory_format=torch.channels_last)\nprint(\"Params (M):\", sum(p.numel() for p in model.parameters())/1e6)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-11T15:32:06.492552Z","iopub.execute_input":"2025-10-11T15:32:06.492860Z","iopub.status.idle":"2025-10-11T15:32:07.670313Z","shell.execute_reply.started":"2025-10-11T15:32:06.492838Z","shell.execute_reply":"2025-10-11T15:32:07.669279Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ===========================\n# Cell 3: Aug utils • Loss • AdamW • Warmup+Cosine • Train/Val Epochs (EPOCHS=200)\n# ===========================\nimport math\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\n# ----- CutMix -----\ndef rand_bbox(W, H, lam):\n    cut_rat = math.sqrt(1. - lam)\n    cut_w, cut_h = int(W * cut_rat), int(H * cut_rat)\n    cx, cy = np.random.randint(W), np.random.randint(H)\n    x1, x2 = np.clip(cx - cut_w // 2, 0, W), np.clip(cx + cut_w // 2, 0, W)\n    y1, y2 = np.clip(cy - cut_h // 2, 0, H), np.clip(cy + cut_h // 2, 0, H)\n    return x1, y1, x2, y2\n\ndef apply_cutmix(x, y, alpha=1.0):\n    if alpha <= 0:\n        return x, y, y, 1.0\n    lam = np.random.beta(alpha, alpha)\n    bs, c, h, w = x.size()\n    index = torch.randperm(bs, device=x.device)\n    x1,y1_,x2,y2_ = rand_bbox(w, h, lam)\n    x[:, :, y1_:y2_, x1:x2] = x[index, :, y1_:y2_, x1:x2]\n    lam = 1 - ((x2-x1)*(y2_-y1_) / (w*h))\n    return x, y, y[index], lam\n\n# ----- Hyperparams / Optim / Scheduler / Loss -----\nEPOCHS        = 200\nbase_lr       = 1.2e-3 * (BATCH / 256)\nweight_decay  = 6e-4\ncutmix_p      = 0.7\ncutmix_alpha  = 1.0\nlabel_smooth  = 0.10\ngrad_clip     = 1.0  # None to disable\n\nopt   = torch.optim.AdamW(model.parameters(), lr=base_lr, weight_decay=weight_decay)\ncrit  = nn.CrossEntropyLoss(label_smoothing=label_smooth)\n\n# --- LR schedule steps must match actual optimizer steps (respect accumulation)\nACCUM = 1 if BATCH >= 256 else 2\nsteps_per_epoch = math.ceil(len(train_dl) / ACCUM)\ntotal_steps  = max(1, EPOCHS * steps_per_epoch)\nwarmup_steps = int(0.10 * total_steps)\n\ndef lr_lambda(step):\n    if step < warmup_steps:\n        return (step + 1) / max(1, warmup_steps)          # linear warmup\n    t = (step - warmup_steps) / max(1, total_steps - warmup_steps)\n    return 0.5 * (1.0 + math.cos(math.pi * t))            # cosine decay\n\nsched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda=lr_lambda)\n\n# AMP scaler (new API)\nscaler = torch.amp.GradScaler('cuda', enabled=(device.type == 'cuda'))\n\n# ----- Train / Val epoch functions -----\ndef run_epoch(dl, train=True, cutmix_prob=cutmix_p, cutmix_alpha=cutmix_alpha):\n    model.train(train)\n    total, correct, total_loss = 0, 0, 0.0\n    step_in_accum = 0\n\n    for x, y in dl:\n        x = x.to(device, non_blocking=True).to(memory_format=torch.channels_last)\n        y = y.to(device, non_blocking=True)\n\n        if train:\n            # Defaults (safe if CutMix is skipped)\n            y1, y2, lam = y, y, 1.0\n            use_cm = (np.random.rand() < cutmix_prob)\n            if use_cm:\n                with torch.no_grad():\n                    # clone avoids in-place mucking with cached non-leaf tensors\n                    x, y1, y2, lam = apply_cutmix(x.clone(), y, alpha=cutmix_alpha)\n\n            with torch.amp.autocast(device_type='cuda', enabled=(device.type == 'cuda')):\n                logits = model(x)\n                loss = lam * crit(logits, y1) + (1 - lam) * crit(logits, y2) if use_cm else crit(logits, y)\n\n            scaler.scale(loss / ACCUM).backward()\n            step_in_accum += 1\n\n            if step_in_accum % ACCUM == 0:\n                if grad_clip is not None:\n                    scaler.unscale_(opt)\n                    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n                scaler.step(opt)\n                scaler.update()\n                opt.zero_grad(set_to_none=True)\n                sched.step()  # step once per optimizer step (aligned with total_steps)\n        else:\n            with torch.no_grad(), torch.amp.autocast(device_type='cuda', enabled=(device.type == 'cuda')):\n                logits = model(x)\n                loss = crit(logits, y)\n\n        total_loss += loss.item() * x.size(0)\n        correct    += (logits.argmax(1) == y).sum().item()\n        total      += x.size(0)\n\n    return total_loss / total, correct / total\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-11T15:32:17.515850Z","iopub.execute_input":"2025-10-11T15:32:17.516670Z","iopub.status.idle":"2025-10-11T15:32:17.537190Z","shell.execute_reply.started":"2025-10-11T15:32:17.516641Z","shell.execute_reply":"2025-10-11T15:32:17.536510Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ===========================\n# Cell 4: Training loop • checkpoint save • optional resume + EMA\n# ===========================\nfrom pathlib import Path\nimport copy\nCKPT_PATH = Path(\"checkpoint.pth\")\nBEST_PATH = Path(\"best.pth\")\nBEST_RAW_PATH = Path(\"best_raw.pth\")   # optional: best non-EMA\nRESUME    = True   # set False to start fresh\n\n# ----- Simple EMA -----\nclass EMA:\n    def __init__(self, model, decay=0.999):\n        self.decay = decay\n        self.shadow = {k: v.detach().clone() for k, v in model.state_dict().items() if v.dtype.is_floating_point}\n    @torch.no_grad()\n    def update(self, model):\n        for k, v in model.state_dict().items():\n            if k in self.shadow and v.dtype.is_floating_point:\n                self.shadow[k].mul_(self.decay).add_(v.detach(), alpha=1.0 - self.decay)\n    def store(self, model):\n        # save current weights to restore later\n        self.backup = {k: v.detach().clone() for k, v in model.state_dict().items()}\n    def copy_to(self, model):\n        model.load_state_dict({**model.state_dict(), **self.shadow}, strict=False)\n    def restore(self, model):\n        if hasattr(self, \"backup\"):\n            model.load_state_dict(self.backup, strict=False)\n            del self.backup\n\nstart_epoch = 1\nbest_acc = 0.0\nbest_raw = 0.0\n\nema = EMA(model, decay=0.999)\n\nif RESUME and CKPT_PATH.exists():\n    print(f\"Resuming from {CKPT_PATH} …\")\n    ckpt = torch.load(CKPT_PATH, map_location=\"cpu\")\n\n    # model -> device (after loading)\n    model.load_state_dict(ckpt[\"model\"], strict=True)\n    model = model.to(device).to(memory_format=torch.channels_last)\n\n    # optimizer / scheduler / scaler\n    opt.load_state_dict(ckpt[\"opt\"])\n    sched.load_state_dict(ckpt[\"sched\"])\n    if \"scaler\" in ckpt:\n        try:\n            scaler.load_state_dict(ckpt[\"scaler\"])\n        except Exception:\n            pass\n\n    # EMA (if present)\n    if \"ema\" in ckpt:\n        ema.shadow = {k: v.to(device if v.is_cuda else \"cpu\") for k, v in ckpt[\"ema\"].items()}\n\n    best_acc    = ckpt.get(\"best_acc\", 0.0)\n    best_raw    = ckpt.get(\"best_raw\", 0.0)\n    start_epoch = ckpt.get(\"epoch\", 0) + 1\n    print(f\"Start at epoch {start_epoch}, best val acc (EMA) {best_acc:.4f}, raw {best_raw:.4f}\")\n\nfor epoch in range(start_epoch, EPOCHS + 1):\n    tr_loss, tr_acc = run_epoch(train_dl, train=True)\n    # update EMA *after* finishing the train epoch (per-step EMA would be better but this still helps)\n    ema.update(model)\n\n    # evaluate raw\n    va_loss_raw, va_acc_raw = run_epoch(val_dl, train=False)\n\n    # evaluate EMA weights\n    ema.store(model)\n    ema.copy_to(model)\n    va_loss_ema, va_acc_ema = run_epoch(val_dl, train=False)\n    ema.restore(model)\n\n    # save latest checkpoint\n    torch.save({\n        \"epoch\": epoch,\n        \"model\": model.state_dict(),\n        \"opt\":   opt.state_dict(),\n        \"sched\": sched.state_dict(),\n        \"scaler\": scaler.state_dict(),\n        \"ema\": ema.shadow,\n        \"best_acc\": best_acc,\n        \"best_raw\": best_raw,\n        \"batch\": BATCH,\n    }, CKPT_PATH)\n\n    # best (raw & EMA)\n    if va_acc_raw > best_raw:\n        best_raw = va_acc_raw\n        torch.save({\"model\": model.state_dict(), \"acc\": best_raw, \"epoch\": epoch}, BEST_RAW_PATH)\n\n    if va_acc_ema > best_acc:\n        best_acc = va_acc_ema\n        # save EMA weights for inference\n        torch.save({\"model\": ema.shadow, \"acc\": best_acc, \"epoch\": epoch}, BEST_PATH)\n\n    print(f\"Epoch {epoch:03d}/{EPOCHS} | \"\n          f\"train {tr_loss:.4f}/{tr_acc:.4f} | \"\n          f\"val_raw {va_loss_raw:.4f}/{va_acc_raw:.4f} | \"\n          f\"val_ema {va_loss_ema:.4f}/{va_acc_ema:.4f} | \"\n          f\"best_raw {best_raw:.4f} | best_ema {best_acc:.4f}\")\n\nprint(f\"Best EMA checkpoint: {BEST_PATH}  | best_ema_acc: {best_acc:.4f}\")\nprint(f\"(Optional) Best raw checkpoint: {BEST_RAW_PATH}  | best_raw_acc: {best_raw:.4f}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-11T15:32:29.644827Z","iopub.execute_input":"2025-10-11T15:32:29.645104Z","iopub.status.idle":"2025-10-11T15:33:50.666739Z","shell.execute_reply.started":"2025-10-11T15:32:29.645085Z","shell.execute_reply":"2025-10-11T15:33:50.665418Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ===========================\n# Cell 5: Inference & Submission (EMA-aware + light TTA)\n# ===========================\nimport torch, pandas as pd\nfrom pathlib import Path\n\nBEST_PATH = Path(\"best.pth\")        # saved with {\"model\": EMA shadow} in Cell 4\nFALLBACK  = Path(\"best_raw.pth\")    # optional raw-best (saved in Cell 4)\n\n# ---- Load best weights (prefer EMA) ----\nckpt_path = BEST_PATH if BEST_PATH.exists() else FALLBACK\nassert ckpt_path.exists(), \"No best checkpoint found. Run training (Cell 4) first.\"\nckpt = torch.load(ckpt_path, map_location=\"cpu\")\nstate = ckpt[\"model\"] if isinstance(ckpt, dict) and \"model\" in ckpt else ckpt\n\n# load and move to device\nmodel.load_state_dict(state, strict=False)\nmodel = model.to(device).to(memory_format=torch.channels_last)\nmodel.eval()\n\n# ---- simple TTA: identity + horizontal flip ----\n@torch.no_grad()\ndef predict_tta(x):\n    with torch.amp.autocast(device_type='cuda', enabled=(device.type=='cuda')):\n        p1 = model(x).softmax(1)\n        p2 = model(torch.flip(x, dims=[3])).softmax(1)\n        return (p1 + p2) * 0.5\n\nif USE_KAGGLE:\n    ids, preds = [], []\n    with torch.no_grad():\n        for x, img_ids in test_dl:\n            x = x.to(device, non_blocking=True).to(memory_format=torch.channels_last)\n            pred = predict_tta(x).argmax(1).tolist()\n            preds.extend([id2label[i] for i in pred])\n            ids.extend(img_ids)\n    sub = pd.DataFrame({\"id\": ids, \"label\": preds}).sort_values(\"id\").reset_index(drop=True)\n    sub.to_csv(\"submission.csv\", index=False)\n    print(f\"Saved submission.csv using {ckpt_path.name}\")\nelse:\n    # Local/debug: report accuracy on torchvision test set\n    test_correct = test_total = 0\n    with torch.no_grad():\n        for x, y in test_dl:\n            x = x.to(device, non_blocking=True).to(memory_format=torch.channels_last)\n            y = y.to(device, non_blocking=True)\n            pred = predict_tta(x).argmax(1)\n            test_correct += (pred == y).sum().item()\n            test_total   += y.numel()\n    print(\"Torchvision CIFAR-10 test accuracy:\", test_correct / test_total)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-11T15:35:45.153366Z","iopub.execute_input":"2025-10-11T15:35:45.153712Z","iopub.status.idle":"2025-10-11T15:41:23.274321Z","shell.execute_reply.started":"2025-10-11T15:35:45.153682Z","shell.execute_reply":"2025-10-11T15:41:23.273171Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}