evaluate_all.py

# ===========================================================
# evaluate_all.py  •  memory-safe, SciPy-free
# ===========================================================
from __future__ import annotations
import re, glob, argparse, math
from pathlib import Path
from typing import List, Dict, Any

import numpy as np
import pandas as pd
import torch, eagerpy as ep
from torchvision import datasets, transforms
from foolbox import PyTorchModel
# --- import only attacks that need no SciPy -----------------
from foolbox import PyTorchModel, attacks, distances

import torch.nn as nn

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ------------------------- models ---------------------------
class MLP(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, depth: int, dropout: float = 0.0):
        super().__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.ModuleList([
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True),
        ])
        for _ in range(depth - 1):
            self.layers.extend([
                nn.Linear(hidden_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.ReLU(inplace=True),
            ])
            if dropout > 0:
                self.layers.append(nn.Dropout(dropout))
        self.out = nn.Linear(hidden_dim, 10)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.flatten(x)
        for layer in self.layers:
            x = layer(x)
        return self.out(x)

class CNN(nn.Module):
    """3-conv CNN; width=k controls channel count."""
    def __init__(self, in_ch: int = 1, k: int = 32, depth: int = 4):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(in_ch,     k,  3, padding=1), nn.BatchNorm2d(k),   nn.ReLU(),
        )
        for _ in range(depth - 1):
            self.layers.extend([
                nn.Conv2d(k,     k,  3, padding=1), nn.BatchNorm2d(k), nn.ReLU(),
            ])
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(k, 10)
        )
    def forward(self, x):
        return self.head(self.layers(x))

ATKS = {
    "FGSM":   attacks.FGSM,
    "DeepFool": attacks.L2DeepFoolAttack,
    "PGD":    attacks.PGD,
    "L2PGD":  attacks.L2PGD
}

def min_l2_batch(fmodel, xb_ep, eps) -> np.ndarray:
    N = len(xb_ep)
    best = np.full(N, np.inf, np.float32)
    # overwrite yb_ep with model predictions
    yb_ep = fmodel(xb_ep).argmax(1)
    yb_ep = ep.astensor(yb_ep).raw
    # run all attacks
    for name, atk_cls in ATKS.items():
        _r, clipped, success = atk_cls()(fmodel, xb_ep, yb_ep, epsilons=eps)
        for ei, adv in enumerate(clipped):
            m = success[ei].raw.bool()
            if not m.any(): continue
            diff = (adv.raw[m] - xb_ep.raw[m]).view(m.sum(), -1)
            d    = diff.norm(p=2, dim=1).cpu().numpy()
            best[m.cpu().numpy()] = np.minimum(best[m.cpu().numpy()], d)
    best[np.isinf(best)] = max(eps)  # attack failed: distance=0
    return best

def count_params(m): return sum(p.numel() for p in m.parameters() if p.requires_grad)

# ------------------------- dataset helper -------------------
def loaders(name: str, batch: int):
    tf = transforms.ToTensor()
    if name == "mnist":
        tr = datasets.MNIST("./data", True ,download=True, transform=tf)
        te = datasets.MNIST("./data", False,download=True, transform=tf)
        in_ch, flat, d_in = 1, 28*28, 784
    elif name == "cifar10":
        tr = datasets.CIFAR10("./data", True, download=True, transform=tf)
        te = datasets.CIFAR10("./data", False, download=True, transform=tf)
        in_ch, flat, d_in = 3, 3*32*32, 1
    else:
        raise ValueError
    full = torch.utils.data.ConcatDataset([tr, te])
    full_loader  = torch.utils.data.DataLoader(full, batch_size=batch,
                    shuffle=False, num_workers=0, pin_memory=True)
    train_loader = torch.utils.data.DataLoader(tr , batch_size=batch,
                    shuffle=False, num_workers=0, pin_memory=True)
    test_loader  = torch.utils.data.DataLoader(te , batch_size=batch,
                    shuffle=False, num_workers=0, pin_memory=True)
    return full_loader, train_loader, test_loader, in_ch, flat, d_in, len(full)

def accuracy(model, loader, flat):
    correct = total = 0
    with torch.no_grad():
        for xb, yb in loader:
            xb = xb.to(DEVICE)
            if flat:
                xb = xb.view(xb.size(0), -1)
            correct += model(xb).argmax(1).cpu().eq(yb).sum().item()
            total   += yb.size(0)
    return 100*correct/total

# ------------------------- main loop ------------------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--glob", default="saved_models/mmlp_s*_data*_w*_d*_*.pt")
    ap.add_argument("--eps", nargs="+", type=float,
                    default=[0.0,0.002,0.01,0.05,0.1])
    ap.add_argument("--batch", type=int, default=256)
    args = ap.parse_args()

    rex = re.compile(r"mmlp_s(?P<seed>\d+)_data(?P<data>\w+)_w(?P<w>\d+)_d(?P<d>\d+)_") # re.compile(r"m(?P<arch>\w+)_s(?P<seed>\d+)_data(?P<data>\w+)_w(?P<w>\d+)_d(?P<d>\d+)_")
    rows: List[Dict[str,Any]] = []

    for fp in sorted(glob.glob(args.glob)):
        m = rex.match(Path(fp).name)
        if not m: continue
        meta = m.groupdict()
        arch, ds, w, d = "mlp", meta["data"], int(meta["w"]), int(meta["d"])

        print(f"\n>>> {Path(fp).name}")

        full_loader, train_loader, test_loader, in_ch, flat, dim_in, N = loaders(ds, args.batch)

        # build & load model ------------------------------------------
        state = torch.load(fp, map_location="cpu")["state_dict"]
        if arch == "mlp":
            model = MLP(flat, w, d).to(DEVICE)
            flat_in = True
        else:
            model = CNN(in_ch=in_ch, k=w, depth=d).to(DEVICE)
            flat_in = False
        model.load_state_dict(state); model.eval()

        acc_tr  = accuracy(model, train_loader, flat_in)
        acc_te  = accuracy(model, test_loader , flat_in)

        fmodel  = PyTorchModel(model, bounds=(0,1))
        S_total = 0.0
        for xb, yb in full_loader:
            xb = xb.to(DEVICE)
            S_total += min_l2_batch(fmodel, ep.astensor(xb), args.eps).sum()
        S_f   = S_total / N
        p_cnt = count_params(model)
        bound = math.sqrt(p_cnt / (N * dim_in))

        rows.append({**meta, "params": p_cnt,
                     "acc_train": acc_tr, "acc_test": acc_te,
                     "S_f": S_f, "bound": bound})

        print(f"{Path(fp).name}: train={acc_tr:.1f}% test={acc_te:.1f}% "
              f"S={S_f:.4f}  bound={bound:.4f}")

    # ------------------- save csvs -----------------------------------
    df = pd.DataFrame(rows)
    df.to_csv("mlp_all_models.csv", index=False)

    summary = (df.astype({"w":int,"d":int,"seed":int})
                 .groupby(["data","w","d"])
                 .agg(acc_train_mean=("acc_train","mean"),
                      acc_test_mean=("acc_test","mean"),
                      S_f_mean=("S_f","mean"),
                      bound_mean=("bound","mean"),
                      params_mean=("params","mean"),
                      acc_train_std=("acc_train","std"),
                      acc_test_std=("acc_test","std"),
                      S_f_std=("S_f","std"),
                      bound_std=("bound","std")))
    summary.to_csv("mlp_summary_by_config.csv")
    print("\n=== mean ± std over seeds ===")
    print(summary.to_string(float_format=lambda x:f"{x:.4f}"))

if __name__ == "__main__":
    main()