|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""faultinj-e4 GPU micro-demo: short DDPM train with invalid membership GT. |
| 3 | +
|
| 4 | +Exploratory only. Labels: faultinj-e4-*. Never writes corrected-s* formal dirs. |
| 5 | +Does not change 0.55 / formal hold / claim ceiling. |
| 6 | +
|
| 7 | +Design (e4-e5-exploratory-design): |
| 8 | + - Train short on FULL CIFAR-10 train (50k) — invalid contract surface |
| 9 | + - Score with H1-like LR on activations? (expensive) OR use simple loss-based |
| 10 | + membership scores as micro instrument for pathology demo |
| 11 | + - Compare to member-only 25k short control |
| 12 | +
|
| 13 | +This micro-demo uses denoising-loss scores as a cheap instrument (not confirmatory H1) |
| 14 | +to ask whether invalid GT + full-train short model manufactures higher AUC than |
| 15 | +member-only short control under held-out fit/score. |
| 16 | +""" |
| 17 | +from __future__ import annotations |
| 18 | + |
| 19 | +import argparse |
| 20 | +import hashlib |
| 21 | +import json |
| 22 | +import sys |
| 23 | +import time |
| 24 | +from pathlib import Path |
| 25 | + |
| 26 | +import numpy as np |
| 27 | +import torch |
| 28 | +import torch.nn.functional as F |
| 29 | +from sklearn.linear_model import LogisticRegression |
| 30 | +from sklearn.metrics import roc_auc_score |
| 31 | +from sklearn.pipeline import make_pipeline |
| 32 | +from sklearn.preprocessing import StandardScaler |
| 33 | +from torch.utils.data import DataLoader, Subset |
| 34 | +from torchvision import datasets, transforms |
| 35 | + |
| 36 | +ROOT = Path(__file__).resolve().parents[2] |
| 37 | +sys.path.insert(0, str(ROOT / "training" / "ddpm-cifar10")) |
| 38 | +from diffusion import GaussianDiffusionTrainer # noqa: E402 |
| 39 | +from model_unet import UNet # noqa: E402 |
| 40 | + |
| 41 | + |
| 42 | +def set_seed(seed: int) -> None: |
| 43 | + np.random.seed(seed) |
| 44 | + torch.manual_seed(seed) |
| 45 | + torch.cuda.manual_seed_all(seed) |
| 46 | + # Avoid cuDNN stream-mismatch crashes seen on longer micro runs (4070 laptop). |
| 47 | + torch.backends.cudnn.benchmark = False |
| 48 | + torch.backends.cudnn.deterministic = True |
| 49 | + try: |
| 50 | + torch.use_deterministic_algorithms(False) |
| 51 | + except Exception: |
| 52 | + pass |
| 53 | + |
| 54 | + |
| 55 | +def build_unet(device: torch.device) -> UNet: |
| 56 | + # slightly smaller for micro budget while staying DDPM-family |
| 57 | + return UNet( |
| 58 | + T=1000, |
| 59 | + ch=64, |
| 60 | + ch_mult=[1, 2, 2, 2], |
| 61 | + attn=[1], |
| 62 | + num_res_blocks=2, |
| 63 | + dropout=0.1, |
| 64 | + ).to(device) |
| 65 | + |
| 66 | + |
| 67 | +def train_short( |
| 68 | + loader: DataLoader, |
| 69 | + steps: int, |
| 70 | + device: torch.device, |
| 71 | + seed: int, |
| 72 | + lr: float = 2e-4, |
| 73 | +) -> UNet: |
| 74 | + set_seed(seed) |
| 75 | + model = build_unet(device) |
| 76 | + trainer = GaussianDiffusionTrainer(model, 1e-4, 0.02, 1000).to(device) |
| 77 | + opt = torch.optim.Adam(model.parameters(), lr=lr) |
| 78 | + model.train() |
| 79 | + it = iter(loader) |
| 80 | + for step in range(1, steps + 1): |
| 81 | + try: |
| 82 | + batch = next(it) |
| 83 | + except StopIteration: |
| 84 | + it = iter(loader) |
| 85 | + batch = next(it) |
| 86 | + x = batch[0] if isinstance(batch, (list, tuple)) else batch |
| 87 | + x = x.to(device, non_blocking=False) |
| 88 | + loss = trainer(x) |
| 89 | + opt.zero_grad(set_to_none=True) |
| 90 | + loss.backward() |
| 91 | + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| 92 | + opt.step() |
| 93 | + if step == 1 or step % max(1, steps // 5) == 0 or step == steps: |
| 94 | + print(f" step {step}/{steps} loss={float(loss):.4f}", flush=True) |
| 95 | + if device.type == "cuda" and step % 200 == 0: |
| 96 | + torch.cuda.synchronize() |
| 97 | + if device.type == "cuda": |
| 98 | + torch.cuda.synchronize() |
| 99 | + return model |
| 100 | + |
| 101 | + |
| 102 | +@torch.no_grad() |
| 103 | +def loss_features( |
| 104 | + model: UNet, |
| 105 | + images: torch.Tensor, |
| 106 | + device: torch.device, |
| 107 | + n_t: int = 4, |
| 108 | + seed: int = 0, |
| 109 | + batch_size: int = 128, |
| 110 | +) -> np.ndarray: |
| 111 | + """Per-image multi-timestep denoising MSE features (cheap instrument).""" |
| 112 | + model.eval() |
| 113 | + trainer = GaussianDiffusionTrainer(model, 1e-4, 0.02, 1000).to(device) |
| 114 | + # CPU generator avoids device-stream coupling issues with cuDNN. |
| 115 | + g = torch.Generator(device="cpu") |
| 116 | + g.manual_seed(seed) |
| 117 | + x_all = images |
| 118 | + N = x_all.size(0) |
| 119 | + ts = torch.linspace(100, 800, n_t).long() |
| 120 | + out = np.zeros((N, n_t), dtype=np.float32) |
| 121 | + for start in range(0, N, batch_size): |
| 122 | + xb = x_all[start : start + batch_size].to(device, non_blocking=False) |
| 123 | + B = xb.size(0) |
| 124 | + for j, t_scalar in enumerate(ts.tolist()): |
| 125 | + t = torch.full((B,), int(t_scalar), device=device, dtype=torch.long) |
| 126 | + noise = torch.randn(xb.shape, generator=g, device="cpu").to(device) |
| 127 | + betas = trainer.betas |
| 128 | + alphas_bar = torch.cumprod(1.0 - betas, dim=0) |
| 129 | + sqrt_ab = torch.sqrt(alphas_bar[t]).view(B, 1, 1, 1).float() |
| 130 | + sqrt_om = torch.sqrt(1.0 - alphas_bar[t]).view(B, 1, 1, 1).float() |
| 131 | + x_t = sqrt_ab * xb + sqrt_om * noise |
| 132 | + pred = model(x_t, t) |
| 133 | + mse = F.mse_loss(pred, noise, reduction="none").mean(dim=(1, 2, 3)) |
| 134 | + out[start : start + B, j] = mse.detach().float().cpu().numpy() |
| 135 | + if device.type == "cuda": |
| 136 | + torch.cuda.synchronize() |
| 137 | + return out |
| 138 | + |
| 139 | + |
| 140 | +def heldout_auc(x_cal, y_cal, x_eval, y_eval) -> float: |
| 141 | + clf = make_pipeline( |
| 142 | + StandardScaler(), |
| 143 | + LogisticRegression(max_iter=2000, solver="lbfgs", random_state=0), |
| 144 | + ) |
| 145 | + clf.fit(x_cal, y_cal) |
| 146 | + scores = clf.decision_function(x_eval) |
| 147 | + # lower loss often => member; flip scores so higher=member if needed by checking direction |
| 148 | + auc = roc_auc_score(y_eval, scores) |
| 149 | + auc_flip = roc_auc_score(y_eval, -scores) |
| 150 | + return float(max(auc, auc_flip)) |
| 151 | + |
| 152 | + |
| 153 | +def main() -> int: |
| 154 | + ap = argparse.ArgumentParser() |
| 155 | + ap.add_argument("--steps", type=int, default=800) |
| 156 | + ap.add_argument("--batch-size", type=int, default=64) |
| 157 | + ap.add_argument("--seed", type=int, default=20260718) |
| 158 | + ap.add_argument("--n-eval", type=int, default=512) |
| 159 | + ap.add_argument( |
| 160 | + "--receipt-name", |
| 161 | + type=str, |
| 162 | + default="faultinj_e4_gpu_micro_receipt.json", |
| 163 | + help="receipt filename under out-dir (avoid overwriting prior step budgets)", |
| 164 | + ) |
| 165 | + ap.add_argument( |
| 166 | + "--dataset-root", |
| 167 | + type=Path, |
| 168 | + default=Path(r"D:/Code/DiffAudit/Download/datasets-readable/cifar10"), |
| 169 | + ) |
| 170 | + ap.add_argument( |
| 171 | + "--out-dir", |
| 172 | + type=Path, |
| 173 | + default=Path("outputs/paper1-corrected-evidence/fault-injection"), |
| 174 | + ) |
| 175 | + args = ap.parse_args() |
| 176 | + |
| 177 | + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 178 | + print(f"device={device} steps={args.steps}", flush=True) |
| 179 | + t0 = time.time() |
| 180 | + set_seed(args.seed) |
| 181 | + |
| 182 | + tfm = transforms.Compose( |
| 183 | + [ |
| 184 | + transforms.ToTensor(), |
| 185 | + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), |
| 186 | + ] |
| 187 | + ) |
| 188 | + # torchvision CIFAR10 root expects parent of cifar-10-batches-py |
| 189 | + ds = datasets.CIFAR10(root=str(args.dataset_root), train=True, download=False, transform=tfm) |
| 190 | + assert len(ds) == 50000 |
| 191 | + |
| 192 | + rng = np.random.default_rng(args.seed) |
| 193 | + all_idx = np.arange(50000) |
| 194 | + member_only = rng.choice(all_idx, size=25000, replace=False) |
| 195 | + member_only.sort() |
| 196 | + |
| 197 | + # invalid GT evaluation rows: sample 512 from train as "members" and 512 as "nonmembers" |
| 198 | + # but model saw ALL 50k. Half of the "nonmembers" are actually trained indices (pathology). |
| 199 | + eval_member = rng.choice(all_idx, size=args.n_eval, replace=False) |
| 200 | + remaining = np.setdiff1d(all_idx, eval_member) |
| 201 | + eval_nonmember_true = rng.choice(remaining, size=args.n_eval // 2, replace=False) |
| 202 | + # false nonmembers: trained indices labeled nonmember |
| 203 | + eval_nonmember_false = rng.choice( |
| 204 | + np.setdiff1d(remaining, eval_nonmember_true), size=args.n_eval // 2, replace=False |
| 205 | + ) |
| 206 | + eval_nonmember = np.concatenate([eval_nonmember_true, eval_nonmember_false]) |
| 207 | + # cal split similarly |
| 208 | + pool = np.setdiff1d(all_idx, np.concatenate([eval_member, eval_nonmember])) |
| 209 | + cal_member = rng.choice(pool, size=args.n_eval, replace=False) |
| 210 | + pool2 = np.setdiff1d(pool, cal_member) |
| 211 | + cal_nm_true = rng.choice(pool2, size=args.n_eval // 2, replace=False) |
| 212 | + cal_nm_false = rng.choice(np.setdiff1d(pool2, cal_nm_true), size=args.n_eval // 2, replace=False) |
| 213 | + cal_nonmember = np.concatenate([cal_nm_true, cal_nm_false]) |
| 214 | + |
| 215 | + def make_loader(indices): |
| 216 | + return DataLoader( |
| 217 | + Subset(ds, indices.tolist()), |
| 218 | + batch_size=args.batch_size, |
| 219 | + shuffle=True, |
| 220 | + num_workers=0, |
| 221 | + drop_last=True, |
| 222 | + ) |
| 223 | + |
| 224 | + print("== train FULL 50k (invalid-contract surface) ==", flush=True) |
| 225 | + model_full = train_short(make_loader(all_idx), args.steps, device, args.seed) |
| 226 | + print("== train MEMBER-ONLY 25k control ==", flush=True) |
| 227 | + model_mo = train_short(make_loader(member_only), args.steps, device, args.seed + 1) |
| 228 | + |
| 229 | + def gather(model, members, nonmembers, seed): |
| 230 | + # load images |
| 231 | + def load_idx(idxs): |
| 232 | + xs = torch.stack([ds[i][0] for i in idxs], dim=0) |
| 233 | + return xs |
| 234 | + |
| 235 | + xm, xn = load_idx(members), load_idx(nonmembers) |
| 236 | + fm = loss_features(model, xm, device, seed=seed) |
| 237 | + fn = loss_features(model, xn, device, seed=seed + 1) |
| 238 | + X = np.concatenate([fm, fn], axis=0) |
| 239 | + y = np.array([1] * len(members) + [0] * len(nonmembers)) |
| 240 | + return X, y |
| 241 | + |
| 242 | + # Invalid-label condition: y uses polluted nonmember definition (includes trained false nonmembers) |
| 243 | + Xc_i, yc_i = gather(model_full, cal_member, cal_nonmember, 11) |
| 244 | + Xe_i, ye_i = gather(model_full, eval_member, eval_nonmember, 12) |
| 245 | + auc_invalid = heldout_auc(Xc_i, yc_i, Xe_i, ye_i) |
| 246 | + |
| 247 | + # Control: member-only model; evaluate with clean held-out nonmembers only (true non-train indices) |
| 248 | + # Build clean nonmembers from indices outside member_only |
| 249 | + non_train = np.setdiff1d(all_idx, member_only) |
| 250 | + cal_m_c = rng.choice(member_only, size=args.n_eval, replace=False) |
| 251 | + cal_n_c = rng.choice(non_train, size=args.n_eval, replace=False) |
| 252 | + eval_m_c = rng.choice(np.setdiff1d(member_only, cal_m_c), size=args.n_eval, replace=False) |
| 253 | + eval_n_c = rng.choice(np.setdiff1d(non_train, cal_n_c), size=args.n_eval, replace=False) |
| 254 | + Xc_c, yc_c = gather(model_mo, cal_m_c, cal_n_c, 21) |
| 255 | + Xe_c, ye_c = gather(model_mo, eval_m_c, eval_n_c, 22) |
| 256 | + auc_control = heldout_auc(Xc_c, yc_c, Xe_c, ye_c) |
| 257 | + |
| 258 | + e4_p3 = bool(auc_invalid >= 0.55 and auc_control < 0.55) |
| 259 | + receipt = { |
| 260 | + "protocol_label": "faultinj-e4-gpu-micro-invalid-gt-2026-07-18", |
| 261 | + "evidence_class": "exploratory_non_confirmatory", |
| 262 | + "run_labels": [ |
| 263 | + "faultinj-e4-invalid-gt-fulltrain-short", |
| 264 | + "faultinj-e4-member-only-control-short", |
| 265 | + ], |
| 266 | + "created_at_unix": int(time.time()), |
| 267 | + "device": str(device), |
| 268 | + "gpu_name": torch.cuda.get_device_name(0) if device.type == "cuda" else None, |
| 269 | + "steps": args.steps, |
| 270 | + "batch_size": args.batch_size, |
| 271 | + "seed": args.seed, |
| 272 | + "instrument": "multi-timestep denoising MSE features + held-out LR (not confirmatory H1)", |
| 273 | + "auc_invalid_gt_fulltrain": auc_invalid, |
| 274 | + "auc_member_only_control": auc_control, |
| 275 | + "delta_invalid_minus_control": auc_invalid - auc_control, |
| 276 | + "E4_P3_pathology_mark": e4_p3, |
| 277 | + "formal_hold": True, |
| 278 | + "claim_ceiling_unchanged": True, |
| 279 | + "elapsed_seconds": round(time.time() - t0, 2), |
| 280 | + } |
| 281 | + raw = json.dumps(receipt, sort_keys=True, separators=(",", ":")).encode() |
| 282 | + receipt["receipt_sha256"] = hashlib.sha256(raw).hexdigest() |
| 283 | + args.out_dir.mkdir(parents=True, exist_ok=True) |
| 284 | + out = args.out_dir / args.receipt_name |
| 285 | + out.write_text(json.dumps(receipt, indent=2) + "\n", encoding="utf-8") |
| 286 | + # also save tiny ckpts optional? skip to save disk |
| 287 | + print(json.dumps({"wrote": str(out).replace("\\", "/"), **{k: receipt[k] for k in [ |
| 288 | + "auc_invalid_gt_fulltrain","auc_member_only_control","E4_P3_pathology_mark","elapsed_seconds","receipt_sha256" |
| 289 | + ]}}, indent=2), flush=True) |
| 290 | + return 0 |
| 291 | + |
| 292 | + |
| 293 | +if __name__ == "__main__": |
| 294 | + raise SystemExit(main()) |
0 commit comments