{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":3649,"databundleVersionId":46718}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, random, time\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.gridspec as gridspec\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader, Subset\nfrom torchvision import datasets, transforms, models","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:16:01.225062Z","iopub.execute_input":"2026-03-01T01:16:01.226008Z","iopub.status.idle":"2026-03-01T01:16:01.230063Z","shell.execute_reply.started":"2026-03-01T01:16:01.225974Z","shell.execute_reply":"2026-03-01T01:16:01.22937Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"CFG = dict(\n    # ── paths ──────────────────────────────────────────────────────────────\n    data_dir   = \"./data\",          # torchvision downloads CIFAR-10 here\n    output_dir = \"./ebm_outputs\",   # all plots + model checkpoint saved here\n\n    # ── dataset ─────────────────────────────────────────────────────────────\n    n_classes  = 10,\n    max_train  = 5_000,    # use None for full 50 k training set\n    max_test   = 500,     # use None for full 10 k test set\n\n    # ── noise ───────────────────────────────────────────────────────────────\n    # Noise is added to simulate corrupted / adversarial inputs.\n    # Images are normalised to [-1, 1], so noise_std=0.3 is significant.\n    noise_type = \"gaussian\",    # \"gaussian\" | \"uniform\" | \"salt_pepper\"\n    noise_std  = 0.3,\n\n    # ── training ────────────────────────────────────────────────────────────\n    epochs      = 30,\n    batch_size  = 128,\n    lr          = 1e-4,\n    # SGLD params used during training to generate negative samples\n    sgld_lr     = 0.01,\n    sgld_steps  = 20,\n    buffer_size = 5_000,\n\n    # ── inference-time purification (Langevin) ──────────────────────────────\n    # More steps = better purification but slower.\n    purif_steps = 40,\n    purif_lr    = 0.01,\n    purif_noise = 0.005,    # small injected noise keeps Langevin stochastic\n\n    # ── visualisation ───────────────────────────────────────────────────────\n    n_grid      = 8,        # number of sample images shown in the grid plot\n)\n\nCLASSES = [\"airplane\",\"automobile\",\"bird\",\"cat\",\"deer\",\n           \"dog\",\"frog\",\"horse\",\"ship\",\"truck\"]\n\nos.makedirs(CFG[\"output_dir\"], exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:17:47.048433Z","iopub.execute_input":"2026-03-01T01:17:47.049035Z","iopub.status.idle":"2026-03-01T01:17:47.054812Z","shell.execute_reply.started":"2026-03-01T01:17:47.049005Z","shell.execute_reply":"2026-03-01T01:17:47.05397Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\ntorch.cuda.manual_seed_all(SEED)\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"[BLOCK 2] Device : {DEVICE}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:17:58.79459Z","iopub.execute_input":"2026-03-01T01:17:58.794884Z","iopub.status.idle":"2026-03-01T01:17:58.801462Z","shell.execute_reply.started":"2026-03-01T01:17:58.794859Z","shell.execute_reply":"2026-03-01T01:17:58.800746Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"MEAN, STD = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)\n\ntrain_tf = transforms.Compose([\n    transforms.RandomHorizontalFlip(),\n    transforms.ToTensor(),\n    transforms.Normalize(MEAN, STD),    # pixel ∈ [-1, 1]\n])\ntest_tf = transforms.Compose([\n    transforms.ToTensor(),\n    transforms.Normalize(MEAN, STD),\n])\n\ntrain_full = datasets.CIFAR10(root=CFG[\"data_dir\"], train=True,\n                               download=True, transform=train_tf)\ntest_full  = datasets.CIFAR10(root=CFG[\"data_dir\"], train=False,\n                               download=True, transform=test_tf)\n\n# Subset for faster experiments\nif CFG[\"max_train\"]:\n    train_ds = Subset(train_full, torch.randperm(len(train_full))[:CFG[\"max_train\"]].tolist())\nelse:\n    train_ds = train_full\n\nif CFG[\"max_test\"]:\n    test_ds = Subset(test_full, torch.randperm(len(test_full))[:CFG[\"max_test\"]].tolist())\nelse:\n    test_ds = test_full\n\ntrain_loader = DataLoader(train_ds, batch_size=CFG[\"batch_size\"],\n                          shuffle=True,  num_workers=2, pin_memory=True)\ntest_loader  = DataLoader(test_ds,  batch_size=CFG[\"batch_size\"],\n                          shuffle=False, num_workers=2, pin_memory=True)\n\nprint(f\"[BLOCK 3] Train : {len(train_ds)} samples\")\nprint(f\"[BLOCK 3] Test  : {len(test_ds)} samples\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:18:20.875842Z","iopub.execute_input":"2026-03-01T01:18:20.87615Z","iopub.status.idle":"2026-03-01T01:18:22.523265Z","shell.execute_reply.started":"2026-03-01T01:18:20.87612Z","shell.execute_reply":"2026-03-01T01:18:22.522553Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class EBMResNet(nn.Module):\n    def __init__(self, n_classes=10):\n        super().__init__()\n        base = models.resnet18(weights=None)\n        base.conv1   = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)\n        base.maxpool = nn.Identity()\n        base.fc      = nn.Linear(512, n_classes)\n        self.net = base\n\n    def forward(self, x):\n        return self.net(x)                          # (B, K) logits\n\n    def energy(self, x):\n        return -self.forward(x).logsumexp(dim=1)    # (B,)  free-energy scalar\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:18:36.805418Z","iopub.execute_input":"2026-03-01T01:18:36.805702Z","iopub.status.idle":"2026-03-01T01:18:36.811525Z","shell.execute_reply.started":"2026-03-01T01:18:36.805677Z","shell.execute_reply":"2026-03-01T01:18:36.810829Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ReplayBuffer:\n    def __init__(self, size, shape=(3, 32, 32)):\n        self.size = size\n        self.buf  = torch.FloatTensor(size, *shape).uniform_(-1, 1)\n\n    def sample(self, n):\n        idx  = torch.randint(0, self.size, (n,))\n        samp = self.buf[idx].clone()\n        mask = torch.rand(n) < 0.05          # 5 % random re-init\n        if mask.any():\n            samp[mask] = torch.FloatTensor(\n                int(mask.sum()), *samp.shape[1:]).uniform_(-1, 1)\n        return samp, idx\n\n    def update(self, s, idx):\n        self.buf[idx] = s.cpu().detach()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:18:46.416963Z","iopub.execute_input":"2026-03-01T01:18:46.417839Z","iopub.status.idle":"2026-03-01T01:18:46.42396Z","shell.execute_reply.started":"2026-03-01T01:18:46.417793Z","shell.execute_reply":"2026-03-01T01:18:46.423203Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def sgld_step(model, x, step_size, noise_scale):\n    x = x.detach().requires_grad_(True)\n    with torch.enable_grad():\n        grad = torch.autograd.grad(model.energy(x).sum(), x)[0]\n    return (x - (step_size / 2) * grad\n              + noise_scale * torch.randn_like(x)).detach().clamp_(-1, 1)\n\ndef run_sgld(model, x_init, steps, step_size, noise_scale=0.01):\n    x = x_init.clone().to(DEVICE)\n    for _ in range(steps):\n        x = sgld_step(model, x, step_size, noise_scale)\n    return x\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:18:58.93074Z","iopub.execute_input":"2026-03-01T01:18:58.931284Z","iopub.status.idle":"2026-03-01T01:18:58.936489Z","shell.execute_reply.started":"2026-03-01T01:18:58.931205Z","shell.execute_reply":"2026-03-01T01:18:58.935774Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def add_noise(x: torch.Tensor, cfg: dict) -> torch.Tensor:\n    t = cfg[\"noise_type\"]\n    if t == \"gaussian\":\n        out = x + cfg[\"noise_std\"] * torch.randn_like(x)\n    elif t == \"uniform\":\n        out = x + cfg[\"noise_std\"] * (2 * torch.rand_like(x) - 1)\n    elif t == \"salt_pepper\":\n        out  = x.clone()\n        mask = torch.rand_like(x) < cfg[\"noise_std\"]\n        vals = torch.where(torch.rand_like(x) > 0.5,\n                           torch.ones_like(x), -torch.ones_like(x))\n        out[mask] = vals[mask]\n    else:\n        raise ValueError(f\"Unknown noise_type: {t!r}\")\n    return out.clamp_(-1, 1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:19:08.748178Z","iopub.execute_input":"2026-03-01T01:19:08.748904Z","iopub.status.idle":"2026-03-01T01:19:08.754029Z","shell.execute_reply.started":"2026-03-01T01:19:08.748871Z","shell.execute_reply":"2026-03-01T01:19:08.753263Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_ebm(model, loader, cfg):\n    model.train()\n    opt   = optim.Adam(model.parameters(), lr=cfg[\"lr\"])\n    sched = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=cfg[\"epochs\"])\n    buf   = ReplayBuffer(cfg[\"buffer_size\"])\n    hist  = {\"loss\": [], \"cls_acc\": []}\n\n    for epoch in range(1, cfg[\"epochs\"] + 1):\n        t0 = time.time()\n        total_loss = correct = n_seen = 0\n\n        for x_real, y in loader:\n            x_real, y = x_real.to(DEVICE), y.to(DEVICE)\n            B = x_real.size(0)\n\n            # generate negative samples\n            x_init, buf_idx = buf.sample(B)\n            x_fake = run_sgld(model, x_init.to(DEVICE),\n                              cfg[\"sgld_steps\"], cfg[\"sgld_lr\"], 0.005)\n            buf.update(x_fake, buf_idx)\n\n            opt.zero_grad()\n            logits = model(x_real)\n            E_r    = model.energy(x_real)\n            E_f    = model.energy(x_fake)\n            loss   = (F.cross_entropy(logits, y)\n                      + (E_r.mean() - E_f.mean())\n                      + 0.01 * (E_r**2 + E_f**2).mean())\n            loss.backward()\n            nn.utils.clip_grad_norm_(model.parameters(), 0.1)\n            opt.step()\n\n            total_loss += loss.item() * B\n            correct    += (logits.argmax(1) == y).sum().item()\n            n_seen     += B\n\n        sched.step()\n        acc = correct / n_seen * 100\n        avg = total_loss / n_seen\n        hist[\"loss\"].append(avg)\n        hist[\"cls_acc\"].append(acc)\n        print(f\"  Epoch {epoch:3d}/{cfg['epochs']} | \"\n              f\"Loss={avg:.4f} | Acc={acc:.1f}% | t={time.time()-t0:.1f}s\")\n\n    return hist\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:20:08.328069Z","iopub.execute_input":"2026-03-01T01:20:08.32859Z","iopub.status.idle":"2026-03-01T01:20:08.336283Z","shell.execute_reply.started":"2026-03-01T01:20:08.328559Z","shell.execute_reply":"2026-03-01T01:20:08.33551Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@torch.no_grad()\ndef evaluate(model, loader, noisy=False, cfg=None):\n    model.eval()\n    Es, Ps, Ls = [], [], []\n    for x, y in loader:\n        x = x.to(DEVICE)\n        if noisy:\n            x = add_noise(x, cfg)\n        Es.append(model.energy(x).cpu().numpy())\n        Ps.append(model(x).argmax(1).cpu().numpy())\n        Ls.append(y.numpy())\n    Es, Ps, Ls = map(np.concatenate, [Es, Ps, Ls])\n    return Es, Ps, Ls, (Ps == Ls).mean() * 100","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:20:18.702808Z","iopub.execute_input":"2026-03-01T01:20:18.703114Z","iopub.status.idle":"2026-03-01T01:20:18.708715Z","shell.execute_reply.started":"2026-03-01T01:20:18.703087Z","shell.execute_reply":"2026-03-01T01:20:18.708015Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def denorm(t):\n    \"\"\"Tensor [-1,1] → numpy [0,1] for imshow.\"\"\"\n    return ((t.permute(1, 2, 0).numpy() + 1) / 2).clip(0, 1)\n\n\ndef plot_4row_grid(clean, noisy, purif,\n                   e_noisy_per, e_purif_per,\n                   labels, cfg, save=None):\n    \"\"\"\n    4-row image grid for N samples side by side.\n\n    Row 0 : Clean original image\n    Row 1 : After noise injection         (energy shown in title)\n    Row 2 : After Langevin purification   (energy shown in title)\n    Row 3 : Bar chart  — energy before vs after purification per image\n    \"\"\"\n    n = len(labels)\n    fig = plt.figure(figsize=(n * 2, 10))\n    gs  = gridspec.GridSpec(4, n, hspace=0.55, wspace=0.08)\n\n    row_labels = [\"① Clean\", f\"② Noisy (σ={cfg['noise_std']})\",\n                  \"③ Purified\", \"④ Energy comparison\"]\n\n    for col in range(n):\n        imgs  = [clean[col], noisy[col], purif[col]]\n        for row in range(3):\n            ax = fig.add_subplot(gs[row, col])\n            ax.imshow(denorm(imgs[row]))\n            ax.axis(\"off\")\n\n            # left-column row labels\n            if col == 0:\n                ax.set_ylabel(row_labels[row], fontsize=9,\n                              labelpad=4, rotation=0, ha=\"right\", va=\"center\")\n\n            # column title: class name (only on row 0)\n            if row == 0:\n                ax.set_title(CLASSES[labels[col]], fontsize=9, pad=3)\n\n            # energy annotations on rows 1 and 2\n            if row == 1:\n                ax.set_xlabel(f\"E={e_noisy_per[col]:.1f}\",\n                              fontsize=7, color=\"red\", labelpad=2)\n            if row == 2:\n                ax.set_xlabel(f\"E={e_purif_per[col]:.1f}\",\n                              fontsize=7, color=\"blue\", labelpad=2)\n\n        # Row 3 — energy bar per image\n        ax_bar = fig.add_subplot(gs[3, col])\n        bars = ax_bar.bar([0, 1],\n                          [e_noisy_per[col], e_purif_per[col]],\n                          color=[\"red\", \"blue\"], width=0.5, alpha=0.8)\n        ax_bar.set_xticks([0, 1])\n        ax_bar.set_xticklabels([\"Noisy\", \"Purif\"], fontsize=7)\n        ax_bar.tick_params(axis=\"y\", labelsize=6)\n        ax_bar.set_title(\"Energy\", fontsize=7, pad=2)\n        ax_bar.grid(axis=\"y\", alpha=0.3, linewidth=0.5)\n        # draw delta\n        delta = e_noisy_per[col] - e_purif_per[col]\n        color = \"green\" if delta > 0 else \"orange\"\n        ax_bar.set_xlabel(f\"Δ={delta:+.1f}\", fontsize=7,\n                          color=color, labelpad=2)\n\n        if col == 0:\n            ax_bar.set_ylabel(row_labels[3], fontsize=9,\n                              labelpad=4, rotation=0, ha=\"right\", va=\"center\")\n\n    fig.suptitle(\n        \"EBM Defense: Clean  →  Noisy  →  Purified\\n\"\n        f\"Noise: {cfg['noise_type']}  σ={cfg['noise_std']}  |  \"\n        f\"Langevin steps: {cfg['purif_steps']}\",\n        fontsize=12, y=1.01\n    )\n\n    if save:\n        plt.savefig(save, bbox_inches=\"tight\", dpi=160)\n    plt.show()\n    print(f\"[PLOT] 4-row grid saved → {save}\")\n\n\ndef plot_energy_dist(Ec, En, Ep, save=None):\n    \"\"\"\n    Overlapping energy histograms for clean / noisy / purified.\n    Expected result: noisy > clean ≈ purified  (purification restores low energy).\n    \"\"\"\n    fig, ax = plt.subplots(figsize=(9, 4))\n    kw = dict(bins=60, alpha=0.65, density=True)\n    ax.hist(Ec, label=f\"Clean    μ={Ec.mean():.2f}  σ={Ec.std():.2f}\",\n            color=\"green\", **kw)\n    ax.hist(En, label=f\"Noisy    μ={En.mean():.2f}  σ={En.std():.2f}\",\n            color=\"red\",   **kw)\n    ax.hist(Ep, label=f\"Purified μ={Ep.mean():.2f}  σ={Ep.std():.2f}\",\n            color=\"royalblue\", **kw)\n\n    # vertical mean lines\n    for val, col in [(Ec.mean(), \"green\"), (En.mean(), \"red\"),\n                     (Ep.mean(), \"royalblue\")]:\n        ax.axvline(val, color=col, linestyle=\"--\", linewidth=1.5)\n\n    ax.set_xlabel(\"E(x) = −log Σ exp f_y(x)\", fontsize=11)\n    ax.set_ylabel(\"Density\")\n    ax.set_title(\"Energy Distribution: Clean vs Noisy vs Purified\", fontsize=13)\n    ax.legend(fontsize=10); ax.grid(alpha=0.3)\n    plt.tight_layout()\n    if save:\n        plt.savefig(save, bbox_inches=\"tight\", dpi=150)\n    plt.show()\n    print(f\"[PLOT] Energy histogram saved → {save}\")\n\n\ndef plot_energy_scatter(e_noisy_per, e_purif_per, save=None):\n    \"\"\"\n    Scatter plot: energy before purification (x) vs after (y).\n    Points below the diagonal = purification reduced energy (good).\n    \"\"\"\n    fig, ax = plt.subplots(figsize=(6, 6))\n    ax.scatter(e_noisy_per, e_purif_per,\n               alpha=0.4, s=15, color=\"royalblue\", edgecolors=\"none\")\n    lims = [min(e_noisy_per.min(), e_purif_per.min()) - 0.5,\n            max(e_noisy_per.max(), e_purif_per.max()) + 0.5]\n    ax.plot(lims, lims, \"r--\", linewidth=1.5, label=\"No change (y=x)\")\n    ax.set_xlim(lims); ax.set_ylim(lims)\n    ax.set_xlabel(\"Energy  BEFORE purification (noisy)\", fontsize=11)\n    ax.set_ylabel(\"Energy  AFTER  purification\", fontsize=11)\n    ax.set_title(\"Per-Sample Energy: Before vs After Purification\", fontsize=12)\n    ax.legend(fontsize=10); ax.grid(alpha=0.3)\n    below = (e_purif_per < e_noisy_per).mean() * 100\n    ax.text(0.05, 0.92, f\"{below:.1f}% of samples have lower energy after purification\",\n            transform=ax.transAxes, fontsize=9, color=\"darkgreen\")\n    plt.tight_layout()\n    if save:\n        plt.savefig(save, bbox_inches=\"tight\", dpi=150)\n    plt.show()\n    print(f\"[PLOT] Energy scatter saved → {save}\")\n\n\ndef plot_training(hist, save=None):\n    fig, (a1, a2) = plt.subplots(1, 2, figsize=(10, 4))\n    ep = range(1, len(hist[\"loss\"]) + 1)\n    a1.plot(ep, hist[\"loss\"],    color=\"steelblue\",  lw=2)\n    a1.set_title(\"Training Loss\"); a1.set_xlabel(\"Epoch\"); a1.grid(alpha=0.3)\n    a2.plot(ep, hist[\"cls_acc\"], color=\"darkorange\", lw=2)\n    a2.set_title(\"Train Accuracy (%)\"); a2.set_xlabel(\"Epoch\"); a2.grid(alpha=0.3)\n    plt.suptitle(\"EBM Training Curves\", fontsize=13); plt.tight_layout()\n    if save:\n        plt.savefig(save, bbox_inches=\"tight\", dpi=150)\n    plt.show()\n    print(f\"[PLOT] Training curves saved → {save}\")\n\n\ndef plot_per_class(labels, pn, pp, save=None):\n    an, ap = [], []\n    for c in range(CFG[\"n_classes\"]):\n        m = labels == c\n        an.append((pn[m] == c).mean() * 100 if m.any() else 0)\n        ap.append((pp[m] == c).mean() * 100 if m.any() else 0)\n    x, w = np.arange(CFG[\"n_classes\"]), 0.35\n    fig, ax = plt.subplots(figsize=(11, 4))\n    ax.bar(x - w/2, an, w, label=\"Noisy\",    color=\"salmon\",     alpha=0.9)\n    ax.bar(x + w/2, ap, w, label=\"Purified\", color=\"steelblue\",  alpha=0.9)\n    ax.set_xticks(x); ax.set_xticklabels(CLASSES, rotation=30, ha=\"right\")\n    ax.set_ylim(0, 100); ax.set_ylabel(\"Accuracy (%)\")\n    ax.set_title(\"Per-Class Classification Accuracy: Noisy vs Purified\", fontsize=13)\n    ax.legend(); ax.grid(axis=\"y\", alpha=0.3)\n    plt.tight_layout()\n    if save:\n        plt.savefig(save, bbox_inches=\"tight\", dpi=150)\n    plt.show()\n    print(f\"[PLOT] Per-class bar chart saved → {save}\")\n\n\ndef print_summary(Ec, ac, En, an, Ep, ap):\n    \"\"\"Formatted comparison table printed to stdout.\"\"\"\n    print(\"\\n\" + \"═\"*65)\n    print(f\"{'  EBM DEFENSE — RESULTS SUMMARY':^65}\")\n    print(\"═\"*65)\n    print(f\"  {'':13}  {'Mean Energy':>13}  {'Std Energy':>11}  {'Acc (%)':>9}\")\n    print(\"─\"*65)\n    for name, E, acc in [(\"① Clean\",    Ec, ac),\n                          (\"② Noisy\",    En, an),\n                          (\"③ Purified\", Ep, ap)]:\n        print(f\"  {name:<13}  {E.mean():>13.3f}  {E.std():>11.3f}  {acc:>9.2f}%\")\n    print(\"═\"*65)\n    print(f\"\\n  Energy drop   (noisy → purified)  : {En.mean()-Ep.mean():>+8.3f}\")\n    print(f\"  Accuracy gain (noisy → purified)  : {ap-an:>+8.2f}%\")\n    print(f\"  Accuracy gap  (clean vs purified)  : {ac-ap:>+8.2f}%\")\n    print(\"═\"*65 + \"\\n\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:21:28.829662Z","iopub.execute_input":"2026-03-01T01:21:28.829931Z","iopub.status.idle":"2026-03-01T01:21:28.85424Z","shell.execute_reply.started":"2026-03-01T01:21:28.829906Z","shell.execute_reply":"2026-03-01T01:21:28.853621Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def main():\n    # ── Instantiate model ─────────────────────────────────────────────────\n    model = EBMResNet(CFG[\"n_classes\"]).to(DEVICE)\n    n_params = sum(p.numel() for p in model.parameters())\n    print(f\"[BLOCK 12] EBMResNet parameters: {n_params:,}\")\n\n    # ── STEP 1: Train ─────────────────────────────────────────────────────\n    print(\"\\n[STEP 1] Training EBM (JEM) …\")\n    hist = train_ebm(model, train_loader, CFG)\n    plot_training(hist, save=f\"{CFG['output_dir']}/training_curves.png\")\n    torch.save(model.state_dict(), f\"{CFG['output_dir']}/ebm_cifar10.pth\")\n    print(f\"[STEP 1] Model saved → {CFG['output_dir']}/ebm_cifar10.pth\")\n\n    # ── STEP 2: Evaluate CLEAN ────────────────────────────────────────────\n    print(\"\\n[STEP 2] Evaluating CLEAN images …\")\n    Ec, _, labels, ac = evaluate(model, test_loader, noisy=False)\n    print(f\"  Clean  → Acc={ac:.2f}%  |  Mean E={Ec.mean():.3f}\")\n\n    # ── STEP 3: Evaluate NOISY ────────────────────────────────────────────\n    print(f\"\\n[STEP 3] Evaluating NOISY images \"\n          f\"(type={CFG['noise_type']}, σ={CFG['noise_std']}) …\")\n    En, pn, _, an = evaluate(model, test_loader, noisy=True, cfg=CFG)\n    print(f\"  Noisy  → Acc={an:.2f}%  |  Mean E={En.mean():.3f}\")\n\n    # ── STEP 4: Purify + Evaluate ─────────────────────────────────────────\n    print(f\"\\n[STEP 4] Running Langevin purification \"\n          f\"({CFG['purif_steps']} steps, lr={CFG['purif_lr']}) …\")\n    purif_imgs, Ep, pp, _ = purify_and_evaluate(model, test_loader, CFG)\n    ap = (pp == labels).mean() * 100\n    print(f\"  Purified → Acc={ap:.2f}%  |  Mean E={Ep.mean():.3f}\")\n\n    # ── STEP 5: Summary table ─────────────────────────────────────────────\n    print_summary(Ec, ac, En, an, Ep, ap)\n\n    # ── STEP 6: Build 8-sample image grid ────────────────────────────────\n    print(\"\\n[STEP 6] Building 4-row image grid (8 samples) …\")\n    rng = np.random.default_rng(SEED)\n    grid_idx = rng.choice(len(test_ds), CFG[\"n_grid\"], replace=False)\n\n    # collect clean images and labels\n    clean_imgs = torch.stack([test_ds[i][0] for i in grid_idx])   # (N,3,32,32)\n    grid_labels = [test_ds[i][1] for i in grid_idx]\n\n    # noisy versions\n    noisy_imgs = add_noise(clean_imgs.to(DEVICE), CFG).cpu()\n\n    # purified versions\n    purif_grid = run_sgld(model,\n                          noisy_imgs.to(DEVICE),\n                          CFG[\"purif_steps\"],\n                          CFG[\"purif_lr\"],\n                          CFG[\"purif_noise\"]).cpu()\n\n    # per-image energies (before and after purification)\n    model.eval()\n    with torch.no_grad():\n        e_noisy_grid = model.energy(noisy_imgs.to(DEVICE)).cpu().numpy()\n        e_purif_grid = model.energy(purif_grid.to(DEVICE)).cpu().numpy()\n\n    plot_4row_grid(\n        clean_imgs, noisy_imgs, purif_grid,\n        e_noisy_grid, e_purif_grid, grid_labels, CFG,\n        save=f\"{CFG['output_dir']}/sample_grid_4row.png\"\n    )\n\n    # ── STEP 7: Population-level plots ───────────────────────────────────\n    print(\"\\n[STEP 7] Population-level energy & accuracy plots …\")\n    plot_energy_dist(Ec, En, Ep,\n        save=f\"{CFG['output_dir']}/energy_distributions.png\")\n\n    plot_energy_scatter(En, Ep,\n        save=f\"{CFG['output_dir']}/energy_scatter.png\")\n\n    plot_per_class(labels, pn, pp,\n        save=f\"{CFG['output_dir']}/per_class_accuracy.png\")\n\n    print(f\"\\n[DONE] All outputs saved to: {CFG['output_dir']}\")\n    return model, hist","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:22:23.12147Z","iopub.execute_input":"2026-03-01T01:22:23.121752Z","iopub.status.idle":"2026-03-01T01:22:23.13272Z","shell.execute_reply.started":"2026-03-01T01:22:23.121729Z","shell.execute_reply":"2026-03-01T01:22:23.132084Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if __name__ == \"__main__\":\n    model, history = main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-01T01:22:33.233572Z","iopub.execute_input":"2026-03-01T01:22:33.233856Z","iopub.status.idle":"2026-03-01T01:58:09.51703Z","shell.execute_reply.started":"2026-03-01T01:22:33.233831Z","shell.execute_reply":"2026-03-01T01:58:09.516307Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}