"""Benchmark Pipelines A (in-house Unet3D), B (nnU-Net v2), C (MONAI UNet) on one GPU: DICE, training time, inference time, model size (params + VRAM). Protocol -------- - Data: identical patient-level splits (data/manifests/split_*.jsonl, from scripts/05_build_splits.py). A and C use the processed 1mm head-cropped percentile-normalized volumes (pimg/plabel); B uses the native source volumes (img/label, joined by key from ntuh.jsonl / m6_labeled.jsonl) — each pipeline's own standard data protocol. - Training: --train-epochs on the first --max-rows train rows (single GPU, bf16 fp-mixing as in the study). Per-pipeline configs (see BENCHMARKING.md): A's DEFAULT is A's own defaults — BatchNorm, ReLU, per-sample 0.5CE+0.5mean-dice loss, AdamW 3e-4 + warmup/cosine, grad clip 5, percentile input, legacy augmentation, 50% fg patches (the 0.5508 @100ep record); A with B's full protocol (instance norm + deep supervision + B's actual DC+CE loss, SGD 1e-2 + PolyLR, B's augmentation, Z-score input, 33% fg patches; 0.5417 @100ep) is --b-aligned; --legacy = A's pre-2026-10-07 config (same as the default but the nnU-style approx nn-loss by default); --b-loss-only = legacy + B's actual loss (ablation). C's DEFAULT is B-aligned (its record config); --legacy reverts C to its pre-2026-10-07 config (BatchNorm, AdamW 3e-4 + warmup/cosine, approx nn-loss by default or --std-loss). "1 epoch" = one pass over the training set under each pipeline's native settings: A/C batch 3, 96^3 random patches; B nnU-Net planned batch (2), planned patch size, SGD 1e-2 PolyLR (250 iters/epoch). B builds an ISOLATED dataset (--nnu-dsid, default 221; the study's Dataset210 is never touched); its one-time raw-build + plan&preprocess wall time is reported separately. - Inference: each pipeline's standard protocol on the held-out test rows: A sliding window 96^3 step 48 + 4-flip TTA; B nnUNetPredictor (gaussian blend, nnU-Net mirroring); C MONAI sliding window (gaussian, 50% overlap) + 4-flip TTA. Mean s/volume over --max-eval-rows (first --eval-warmup cases untimed) + peak VRAM of the inference run. - Dice: tumor probability thresholded at 0.5 vs ground truth — the primary metric convention of all three study pipelines (B also reports the argmax hard-seg dice it exports natively). - VRAM training: max allocation over the training run (A/C in-process via torch; B via nvidia-smi polling around the nnUNetv2_train subprocess). Usage (repo root, inside the `longitudinal-313` conda env): python scripts/benchmark_pipelines.py # train + eval A B C python scripts/benchmark_pipelines.py --pipelines A C python scripts/benchmark_pipelines.py --train-epochs 3 --max-rows 240 python scripts/benchmark_pipelines.py --pipelines "A B C" --resume --train-epochs 200 # ^ continue from runs_bench/{A,C}/best.pt and the nnU-Net fold_0 # checkpoint_latest.pth; --train-epochs is the TOTAL budget python scripts/benchmark_pipelines.py --no-tta # no test-time augmentation python scripts/benchmark_pipelines.py --skip-train # eval existing study checkpoints Outputs: results/benchmark_pipelines/benchmark_pipelines.json (+ _per_row.jsonl, _bench.png) """ import argparse import gc import json import math import os import shutil import subprocess import sys import threading import time import traceback from datetime import datetime ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, ROOT) sys.path.insert(0, os.path.join(ROOT, "scripts_nnu")) import numpy as np import torch from src.common import DATA, load_jsonl, read_nii_arr from src.training import dice_np def d(name): p = os.path.join(ROOT, name) os.makedirs(p, exist_ok=True) return p def parse_args(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--pipelines", default="A B C", help="space-separated subset of 'A B C'") ap.add_argument("--gpu", type=int, default=0, help="physical GPU index (all pipelines use one GPU)") ap.add_argument("--train-epochs", type=int, default=100, help="training budget for A and C (and B via NNU_PL_EPOCHS)") ap.add_argument("--max-rows", type=int, default=None, help="train rows for A/C (first N of labeled split_train; default 80%% of the split)") ap.add_argument("--nnu-cases", type=int, default=None, help="train cases for B (first N labeled native cases; default same as --max-rows)") ap.add_argument("--max-eval-rows", type=int, default=None, help="held-out test volumes for inference + DICE (default 10%% of labeled split_test)") ap.add_argument("--eval-warmup", type=int, default=1, help="untimed warmup cases before inference timing") ap.add_argument("--workers", type=int, default=4, help="dataloader workers (A/C)") ap.add_argument("--batch", type=int, default=3, help="per-GPU batch (A/C)") ap.add_argument("--patch", type=int, default=96, help="patch/window (A/C)") ap.add_argument("--a-fg-dice", action="store_true", help="A: use tumor-class-only soft dice in the per-sample loss " "(instead of the background+tumor mean)") ap.add_argument("--a-fg-ce-weight", type=float, default=1.0, help="A: tumor class weight in the CE part of the loss (1.0 = unweighted)") ap.add_argument("--eval-at", default=None, help="A/C: space-separated epochs (1-based) at which to run a quick test eval " "during training, e.g. '5 10' — results land in res['evals']") ap.add_argument("--a-leaky", action="store_true", help="A: LeakyReLU(0.01) instead of ReLU in Unet3D (matches Pipeline C's activation)") ap.add_argument("--nn-loss", action="store_true", default=None, help="A/C: nnU-style approx DC+CE loss (unweighted CE + batch-level soft " "dice over both classes, smooth 1.0, 1:1). The DEFAULT loss in " "--legacy mode (A and C); A's default mode uses --std-loss instead. " "Ignored in b-aligned / b-loss-only modes. A's --a-fg-* flags are " "ignored with this loss") ap.add_argument("--std-loss", action="store_true", help="A/C: per-sample 0.5CE+0.5dice loss. A's DEFAULT mode uses it; in " "--legacy mode it opts out of the approx nn-loss (for A that " "reproduces A's current default mode)") ap.add_argument("--b-aligned", action="store_true", help="A: train with B's full protocol — nnU-Net v2 3d_fullres as read from " "the installed nnunetv2: InstanceNorm3d(affine) + kaiming-normal(0.01) " "init, B's ACTUAL loss (CE + tumor-only batch soft dice, smooth 1e-5) " "with deep supervision (1x/1/2/1/4, weights 4/7, 2/7, 1/7), SGD 1e-2 " "(momentum 0.99, nesterov, wd 3e-5) + PolyLR (exponent 0.9 over the " "total budget), B's augmentation, per-volume Z-score input, 33%% " "foreground patch sampling, grad clip 12, B-style val pseudo-dice EMA " "(0.9/0.1) -> runs_bench/A/best_ema.pt, A uses LeakyReLU. Opt-in for A " "since 2026-10-07 (its record config is the A-native default); still " "the DEFAULT for C (0.4002 @100ep). A 0.5417 @100ep (see " "BENCHMARKING.md, B-aligned A/C runs)") ap.add_argument("--legacy", action="store_true", help="A/C: the pre-2026-10-07 config (BatchNorm, ReLU for A, AdamW 3e-4 + " "warmup/cosine, legacy augmentation, percentile [0,1] input, no deep " "supervision, no val-EMA); loss defaults to the nnU-style approx " "nn-loss (--std-loss opts out — for A that is the current default " "mode). Opt out of B-aligned: A's default is A-native (0.5508 " "@100ep), C's default is B-aligned (0.4002 @100ep)") ap.add_argument("--b-loss-only", action="store_true", help="A/C: legacy config but with B's ACTUAL loss function (b_dc_and_ce: " "CE + tumor-only batch soft dice, smooth 1e-5, do_bg=False) instead of " "the approximate nn-loss. Ablation of the default B-aligned config " "(which showed the loss alone does not fix C's instability)") ap.add_argument("--no-tta", action="store_true", help="disable each pipeline's standard TTA") ap.add_argument("--skip-train", action="store_true", help="eval existing study checkpoints (runs/round0, runs_monai/round0, nnu/results 210) " "instead of training") ap.add_argument("--resume", action="store_true", help="continue training from the latest checkpoint of the SAME config " "mode. A: runs_bench/A/best.pt (default A-native config) / best_bal.pt " "(--b-aligned) / best_legacy.pt (--legacy) / best_bloss.pt " "(--b-loss-only); C: best.pt (default B-aligned) / best_legacy.pt " "(--legacy) / best_bloss.pt. (Model + optimizer + scheduler when " "present.) B: append nnUNet's --c flag (continue from " "fold_0/checkpoint_latest.pth; pair with --nnu-skip-prep). " "--train-epochs is the TOTAL budget, not epochs added") ap.add_argument("--nnu-dsid", type=int, default=None, help="nnU-Net dataset id: training mode default 221 (isolated; the study's 210 is never " "touched), --skip-train mode default 210 (study model); pass explicitly to evaluate " "a previously benchmarked model folder") ap.add_argument("--nnu-skip-prep", action="store_true", help="B/train: reuse an existing raw dataset + " "preprocessed tree for --nnu-dsid") ap.add_argument("--nnu-ckpt", default="checkpoint_best.pth", help="B: checkpoint file inside fold_0 used for the final test inference " "(default checkpoint_best.pth; use checkpoint_final.pth to score the " "model at the last trained epoch — the A/C-comparable trajectory point)") ap.add_argument("--out", default=None, help="output json path (default results/benchmark_pipelines/benchmark_pipelines.json)") return ap.parse_args() # ---------------- data selection ---------------- def native_manifests(): """key -> {img, label} from the labeled source manifests (ntuh + m6_labeled).""" man = os.path.join(DATA, "manifests") out = {} for name in ("ntuh.jsonl", "m6_labeled.jsonl"): p = os.path.join(man, name) if not os.path.exists(p): continue for r in load_jsonl(p): if r.get("label"): out[r["key"]] = {"img": r["img"], "label": r["label"]} return out def native_rows_for(keys): nm = native_manifests() rows = [] for k in keys: if k not in nm: raise SystemExit(f"B: key {k[:60]}... not in ntuh/m6_labeled manifests; " f"native img/label unavailable") rows.append({"key": k, **nm[k]}) return rows def select_rows(args): man = os.path.join(DATA, "manifests") train = [r for r in load_jsonl(os.path.join(man, "split_train.jsonl")) if r.get("plabel")] test = [r for r in load_jsonl(os.path.join(man, "split_test.jsonl")) if r.get("plabel")] val_f = os.path.join(man, "split_val.jsonl") val = [r for r in load_jsonl(val_f) if r.get("plabel")] if os.path.exists(val_f) else [] if not train: raise SystemExit("split_train.jsonl has no labeled processed rows; run scripts/preprocess.py + 05_build_splits.py") if not test: raise SystemExit("split_test.jsonl has no labeled processed rows; run scripts/preprocess.py + 05_build_splits.py") if args.max_rows is None: args.max_rows = max(1, round(0.8 * len(train))) if args.max_eval_rows is None: args.max_eval_rows = max(1, round(0.1 * len(test))) if args.nnu_cases is None: args.nnu_cases = args.max_rows return train[:args.max_rows], test[:args.max_eval_rows], val def val_pseudo_dice(model, val_dl, dev, n_batches=50): """B-style per-epoch validation pseudo-dice: mean per-patch tumor soft dice over the first n_batches of unaugmented val patches (B: 50 val iterations, MemoryEfficientSoftDiceLoss on the foreground class). Accepts A's tuple batches or C's dict batches.""" import torch.nn.functional as F model.eval() vals = [] with torch.no_grad(): for i, item in enumerate(val_dl): if i >= n_batches: break if isinstance(item, dict): img, lab = item["pimg"], item["plabel"].squeeze(1) else: img, lab = item[0], item[1] img = img.to(dev, non_blocking=True) lab = lab.to(dev, non_blocking=True) with torch.autocast("cuda", dtype=torch.bfloat16): logits = model(img) if isinstance(logits, (tuple, list)): logits = logits[0] p = F.softmax(logits.float(), dim=1)[0, 1] t = (lab > 0).float() den = p.sum() + t.sum() vals.append(float((2 * (p * t).sum() + 1) / (den + 1))) model.train() return float(np.mean(vals)) if vals else 0.0 # ---------------- GPU / timing helpers ---------------- def reset_peak(dev): torch.cuda.reset_peak_memory_stats(dev) def peak_mb(dev): return round(torch.cuda.max_memory_allocated(dev) / 1e6, 1) class SmiPoller(threading.Thread): """Samples nvidia-smi memory.used (MiB) for one physical GPU in the background.""" def __init__(self, gpu): super().__init__(daemon=True) self.gpu, self._stop = gpu, False self.max_mib = 0.0 def run(self): while not self._stop: try: out = subprocess.run(["nvidia-smi", f"--id={self.gpu}", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=5).stdout.strip() self.max_mib = max(self.max_mib, float(out.splitlines()[0])) except Exception: pass time.sleep(0.4) def stop(self): self._stop = True self.join(timeout=2) def free_gpu(): gc.collect() torch.cuda.empty_cache() def load_native_label(lab_path, img_itk): lab = read_nii_arr(lab_path).astype("uint8") import SimpleITK as sitk if lab.shape != tuple(img_itk.GetSize()[::-1]): lab = sitk.GetArrayFromImage(sitk.Resample(sitk.ReadImage(lab_path), img_itk, sitk.Transform(), sitk.sitkNearestNeighbor, 0.0)) lab = lab.astype("uint8") return lab # ---------------- Pipeline A ---------------- def bench_A(args, train_rows, test_rows, val_rows=None): from src.dataset import make_dataloader from src.losses import (per_sample_loss, one_hot, nn_dc_and_ce, b_dc_and_ce, deep_sup_loss) from src import training as A import torch.nn.functional as TF dev = torch.device(f"cuda:{args.gpu}") torch.cuda.set_device(dev) ba = args.a_ba tta = not args.no_tta a_default = (not ba) and (not args.legacy) and (not args.b_loss_only) and (not args.a_nn_loss) res = {"pipeline": "A", "backbone": ("Unet3D(base=16, depth=4, leaky, instance-norm, deep-sup) [B-aligned]" if ba else f"Unet3D(base=16, depth=4, {'leaky' if args.a_leaky else 'relu'})" + (" [A default]" if a_default else "")), "b_aligned": ba, "tta": tta, "n_train": len(train_rows), "n_eval": len(test_rows), "loss": ("B's DC_and_CE (CE + tumor-only batch dice, smooth 1e-5) + deep sup " "(1x/1/2/1/4, weights 4/7, 2/7, 1/7)" if ba else ("b_dc_and_ce ONLY (CE + tumor-only batch dice, smooth 1e-5; rest legacy)" if args.b_loss_only else ("nn_dc_and_ce (1:1, batch dice, do_bg, no CE class weight)" if args.a_nn_loss else f"per-sample 0.5CE+0.5mean-dice (fg_only={args.a_fg_dice}, " f"ce_w={args.a_fg_ce_weight})"))), "optimizer": ("SGD 1e-2 (momentum 0.99, nesterov, wd 3e-5) + PolyLR(exp 0.9)" if ba else "AdamW 3e-4 + warmup/cosine")} if ba: res["input"] = "per-volume Z-score; B augmentation; 33% fg patches" res["epochs_total"] = args.train_epochs eval_at = {int(x) for x in (args.eval_at or "").replace(",", " ").split() if x.strip()} def run_eval(m, tag, timed=True): """Test-set dice eval (mid-training or final); results in res['evals'][tag].""" m.eval() step = max(args.patch // 2, 16) times, dices, pmaxs = [], [], [] with torch.no_grad(): for r in test_rows[args.eval_warmup:]: vol = read_nii_arr(r["pimg"]).astype("float32") torch.cuda.synchronize(dev) t0 = time.time() probs = A.sliding_window_probs(m, vol, dev, args.patch, step, tta, zscore=ba) torch.cuda.synchronize(dev) times.append(time.time() - t0) osh = vol.shape probs = probs[:osh[0], :osh[1], :osh[2]] pmaxs.append(float(probs.max())) lab = read_nii_arr(r["plabel"]).astype("uint8") if lab.shape != osh: lab = np.pad(lab, [(0, max(osh[i] - lab.shape[i], 0)) for i in range(3)])[:osh] dices.append(dice_np(probs, lab)) out = {"dice": round(float(np.mean(dices)), 4), "rows_gt_0.001": int(sum(x > 0.001 for x in dices)), "pmax_median": round(float(np.median(pmaxs)), 4)} if timed: out["infer_s_per_vol"] = round(float(np.mean(times)), 2) res.setdefault("evals", {})[tag] = out print(f"[A] eval {tag}: " + " ".join(f"{k}={v}" for k, v in out.items()), flush=True) return dices, pmaxs, times ckpt_dir = d("runs_bench/A") # A's default (A-native all-defaults) writes best.pt; the other modes keep # separate names so the record checkpoint is never clobbered if ba: ckpt_name = "best_bal.pt" elif args.b_loss_only: ckpt_name = "best_bloss.pt" elif args.legacy or args.a_nn_loss: ckpt_name = "best_legacy.pt" elif args.a_leaky or args.a_fg_dice or args.a_fg_ce_weight != 1.0: ckpt_name = "best_alt.pt" else: ckpt_name = "best.pt" ckpt = os.path.join(ckpt_dir, ckpt_name) best_ema_ckpt = os.path.join(ckpt_dir, "best_ema.pt") if args.skip_train: ckpt = os.path.join(ROOT, "runs", "round0", "best.pt") if not os.path.exists(ckpt): raise FileNotFoundError(f"--skip-train: {ckpt} not found") model, _ = A.load_model(ckpt, dev) res["ckpt"] = ckpt else: if ba: from src.unet3d import Unet3D model = Unet3D(in_ch=1, out_ch=2, base=16, depth=4, leaky=True, norm="instance", deep_sup=True, b_init=True).to("cuda") else: model = A.build_model(base=16, device="cuda", leaky=args.a_leaky) patch = (args.patch, args.patch, args.patch) _, dl = make_dataloader(train_rows, patch, args.batch, True, num_workers=args.workers, seed=0, persistent=False, pos_neg=0.33 if ba else 0.5, zscore=ba, b_aug=ba) steps = max(len(dl), 1) if ba: opt = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.99, nesterov=True, weight_decay=3e-5) sched = None else: total_steps = steps * args.train_epochs opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4) warmup = min(300, max(10, total_steps // 10)) sched = torch.optim.lr_scheduler.SequentialLR( opt, [torch.optim.lr_scheduler.LinearLR(opt, start_factor=0.1, total_iters=warmup), torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1), eta_min=3e-4 * 0.05)], milestones=[warmup]) start_epoch = 0 ema = None best_ema = None val_dl = None if ba and val_rows: _, val_dl = make_dataloader(val_rows[:275], patch, 2, True, num_workers=args.workers, seed=1, persistent=True, pos_neg=0.33, zscore=True, no_aug=True) if args.resume and os.path.exists(ckpt): sd = torch.load(ckpt, map_location=dev) model.load_state_dict(sd["model"]) if "optimizer" in sd: opt.load_state_dict(sd["optimizer"]) if "scheduler" in sd and sched is not None: sched.load_state_dict(sd["scheduler"]) start_epoch = int(sd.get("epoch", 0)) res["resumed_from_epoch"] = start_epoch print(f"[A] resume: {ckpt} epoch {start_epoch}/{args.train_epochs} " f"(optimizer={'restored' if 'optimizer' in sd else 'fresh'}, " f"scheduler={'restored' if 'scheduler' in sd else 'fresh'})", flush=True) elif args.resume: print(f"[A] --resume: {ckpt} not found, training from scratch", flush=True) model.train() reset_peak(dev) t0 = time.time() for ep in range(start_epoch, args.train_epochs): if hasattr(dl.sampler, "set_epoch"): dl.sampler.set_epoch(ep) if ba: # PolyLR (B): lr = 1e-2 * (1 - ep/total)^0.9 on the GLOBAL epoch index lr0 = 1e-2 * (1 - ep / args.train_epochs) ** 0.9 for pg in opt.param_groups: pg["lr"] = lr0 else: lr0 = opt.param_groups[0]["lr"] loss_sum = n = 0.0 t_ep = time.time() for img, lab, wts in dl: img = img.to(dev, non_blocking=True) lab = lab.to(dev, non_blocking=True) wts = wts.to(dev, non_blocking=True) with torch.autocast("cuda", dtype=torch.bfloat16): logits = model(img) if ba: l2 = TF.interpolate(lab.unsqueeze(1).float(), scale_factor=0.5, mode="nearest").squeeze(1).long() l4 = TF.interpolate(lab.unsqueeze(1).float(), scale_factor=0.25, mode="nearest").squeeze(1).long() loss = deep_sup_loss(logits, [lab, l2, l4], b_dc_and_ce) loss_val = float(loss.detach()) if not math.isfinite(loss_val): loss = torch.zeros_like(loss) loss_val = 0.0 elif args.b_loss_only: loss = b_dc_and_ce(logits.float(), lab) loss_val = float(loss.detach()) if not math.isfinite(loss_val): loss = torch.zeros_like(loss) loss_val = 0.0 elif args.a_nn_loss: loss = nn_dc_and_ce(logits.float(), lab) loss_val = float(loss.detach()) if not math.isfinite(loss_val): loss = torch.zeros_like(loss) loss_val = 0.0 else: losses = per_sample_loss(logits.float(), one_hot(lab, 2), fg_only=args.a_fg_dice, fg_ce_weight=args.a_fg_ce_weight) if torch.isnan(losses).any() or torch.isinf(losses).any(): losses = torch.zeros_like(losses) loss = (losses * wts).sum() / wts.sum().clamp(min=1e-6) loss_val = float(losses.detach().mean()) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 12.0 if ba else 5.0) opt.step() if sched is not None: sched.step() loss_sum += loss_val n += 1 res.setdefault("per_epoch_loss", []).append(round(loss_sum / max(n, 1), 4)) res.setdefault("per_epoch_time", []).append(round(time.time() - t_ep, 1)) res.setdefault("per_epoch_lr", []).append(round(lr0, 8)) line = (f"[A] epoch {ep + 1}/{args.train_epochs} loss={loss_sum / max(n, 1):.4f} " f"t={time.time() - t_ep:.0f}s lr={lr0:.2e}") if ba and val_dl is not None: vd = val_pseudo_dice(model, val_dl, dev) ema = vd if ema is None else 0.9 * ema + 0.1 * vd res.setdefault("per_epoch_val_dice", []).append(round(vd, 4)) res.setdefault("per_epoch_val_ema", []).append(round(ema, 4)) if best_ema is None or ema > best_ema: best_ema = ema torch.save({"model": model.state_dict(), "epoch": ep + 1, "ema": round(ema, 4)}, best_ema_ckpt) line += f" val={vd:.4f} ema={ema:.4f} best_ema={best_ema:.4f}" print(line, flush=True) if (ep + 1) in eval_at: print(f"[A] mid-eval @ epoch {ep + 1}", flush=True) run_eval(model, f"epoch_{ep + 1}") model.train() if (ep + 1) in eval_at or (ep + 1) % 10 == 0: _dump_partial(res, args) dt = time.time() - t0 res["train_s"] = round(dt, 1) res["train_s_per_epoch"] = round(dt / max(args.train_epochs - start_epoch, 1), 1) res["train_vram_mb"] = peak_mb(dev) res["steps_per_epoch"] = steps if ba: res["best_ema"] = round(best_ema, 4) if best_ema is not None else None res["best_ema_ckpt"] = best_ema_ckpt torch.save({"model": model.state_dict(), "epoch": args.train_epochs, "optimizer": opt.state_dict()}, ckpt) else: torch.save({"model": model.state_dict(), "epoch": args.train_epochs, "optimizer": opt.state_dict(), "scheduler": sched.state_dict()}, ckpt) del dl, opt if sched is not None: del sched res["ckpt"] = res.get("ckpt", ckpt) res["params"] = int(sum(p.numel() for p in model.parameters())) res["params_M"] = round(res["params"] / 1e6, 3) res["weight_mb"] = round(os.path.getsize(res["ckpt"]) / 1e6, 3) # inference: sliding window, 4-flip TTA (study protocol of 06/08) model.eval() step = max(args.patch // 2, 16) warm = test_rows[0] with torch.no_grad(): A.sliding_window_probs(model, read_nii_arr(warm["pimg"]).astype("float32"), dev, args.patch, step, tta, zscore=ba) torch.cuda.synchronize(dev) reset_peak(dev) dices, pmaxs, times = run_eval(model, "final", timed=True) res["infer_s_per_vol"] = round(float(np.mean(times)), 2) if times else None res["infer_s_per_vol_std"] = round(float(np.std(times)), 2) if times else None res["infer_vram_mb"] = peak_mb(dev) res["dice"] = round(float(np.mean(dices)), 4) if dices else 0.0 res["dice_per_row"] = [{"key": test_rows[args.eval_warmup + i]["key"], "dice": round(dices[i], 4), "s": round(times[i], 2), "pmax": round(float(pm), 4)} for i, pm in zip(range(len(dices)), pmaxs)] del model free_gpu() return res # ---------------- Pipeline C (MONAI) ---------------- def bench_C(args, train_rows, test_rows, val_rows=None): sys.path.insert(0, os.path.join(ROOT, "scripts_monai")) from monai_common import (build_model, WeightedDiceCELoss, NNDCeLoss, make_dataloader, predict_probs, prob_dice, load_model, train_transform, rows_for_dataset, _collate as monai_collate) from src.losses import b_dc_and_ce, deep_sup_loss import torch.nn.functional as TF dev = torch.device(f"cuda:{args.gpu}") torch.cuda.set_device(dev) ba = args.c_ba tta = not args.no_tta res = {"pipeline": "C", "backbone": ("monai UNet(16->128, instance-norm, deep-sup) [B-aligned]" if ba else "monai UNet(16->128)"), "b_aligned": ba, "tta": tta, "n_train": len(train_rows), "n_eval": len(test_rows), "loss": ("B's DC_and_CE (CE + tumor-only batch dice, smooth 1e-5) + deep sup " "(1x/1/2/1/4, weights 4/7, 2/7, 1/7)" if ba else ("b_dc_and_ce ONLY (CE + tumor-only batch dice, smooth 1e-5; rest legacy)" if args.b_loss_only else "nn_dc_ce (1:1, batch dice, do_bg, no CE class weight)" if args.c_nn_loss else "per-sample 0.5CE+0.5mean-dice (WeightedDiceCELoss)")), "optimizer": ("SGD 1e-2 (momentum 0.99, nesterov, wd 3e-5) + PolyLR(exp 0.9)" if ba else "AdamW 3e-4 + warmup/cosine")} if ba: res["input"] = "per-volume Z-score; B augmentation; 33% fg patches" res["epochs_total"] = args.train_epochs eval_at = {int(x) for x in (args.eval_at or "").replace(",", " ").split() if x.strip()} def run_eval(m, tag, timed=True): """Test-set dice eval (mid-training or final); results in res['evals'][tag].""" m.eval() times, dices = [], [] with torch.no_grad(): for r in test_rows[args.eval_warmup:]: vol = read_nii_arr(r["pimg"]).astype("float32") torch.cuda.synchronize(dev) t0 = time.time() p = predict_probs(m, vol, dev, args.patch, 0.5, tta=tta, sw_batch=8, zscore=ba) torch.cuda.synchronize(dev) times.append(time.time() - t0) lab = read_nii_arr(r["plabel"]).astype("uint8") if p.shape != lab.shape: lab = np.pad(lab, [(0, max(p.shape[i] - lab.shape[i], 0)) for i in range(3)])[:p.shape] dices.append(prob_dice(p, lab)) out = {"dice": round(float(np.mean(dices)), 4), "rows_gt_0.001": int(sum(x > 0.001 for x in dices))} if timed: out["infer_s_per_vol"] = round(float(np.mean(times)), 2) res.setdefault("evals", {})[tag] = out print(f"[C] eval {tag}: " + " ".join(f"{k}={v}" for k, v in out.items()), flush=True) return dices, times ckpt_dir = d("runs_bench/C") # B-aligned is the default (writes best.pt / best_ema.pt); legacy and # b-loss-only ablation runs keep separate names so the default is never clobbered ckpt_name = "best.pt" if ba else ("best_bloss.pt" if args.b_loss_only else "best_legacy.pt") ckpt = os.path.join(ckpt_dir, ckpt_name) best_ema_ckpt = os.path.join(ckpt_dir, "best_ema.pt") if args.skip_train: ckpt = os.path.join(ROOT, "runs_monai", "round0", "best.pt") if not os.path.exists(ckpt): raise FileNotFoundError(f"--skip-train: {ckpt} not found") model, _ = load_model(ckpt, dev) res["ckpt"] = ckpt else: model = build_model(dev, norm="instance" if ba else "batch", b_aligned=ba) dl = make_dataloader(train_rows, win=args.patch, batch=args.batch, workers=args.workers, b_aligned=ba) steps = max(len(dl), 1) if ba: opt = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.99, nesterov=True, weight_decay=3e-5) sched = None else: total_steps = steps * args.train_epochs opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4) warmup = min(300, max(10, total_steps // 10)) sched = torch.optim.lr_scheduler.SequentialLR( opt, [torch.optim.lr_scheduler.LinearLR(opt, start_factor=0.1, total_iters=warmup), torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1), eta_min=3e-4 * 0.05)], milestones=[warmup]) if ba or args.b_loss_only: loss_f = b_dc_and_ce else: loss_f = NNDCeLoss() if args.c_nn_loss else WeightedDiceCELoss() start_epoch = 0 ema = None best_ema = None val_dl = None if ba and val_rows: from monai.data import Dataset as MDS val_ds = MDS(data=rows_for_dataset(val_rows[:275]), transform=train_transform(args.patch, b_aligned=True, no_aug=True)) val_dl = torch.utils.data.DataLoader(val_ds, batch_size=2, shuffle=True, num_workers=args.workers, collate_fn=monai_collate, drop_last=False, pin_memory=True, persistent_workers=False) if args.resume and os.path.exists(ckpt): sd = torch.load(ckpt, map_location=dev) model.load_state_dict(sd["model"]) if "optimizer" in sd: opt.load_state_dict(sd["optimizer"]) if "scheduler" in sd and sched is not None: sched.load_state_dict(sd["scheduler"]) start_epoch = int(sd.get("epoch", 0)) res["resumed_from_epoch"] = start_epoch print(f"[C] resume: {ckpt} epoch {start_epoch}/{args.train_epochs} " f"(optimizer={'restored' if 'optimizer' in sd else 'fresh'}, " f"scheduler={'restored' if 'scheduler' in sd else 'fresh'})", flush=True) elif args.resume: print(f"[C] --resume: {ckpt} not found, training from scratch", flush=True) model.train() reset_peak(dev) t0 = time.time() for ep in range(start_epoch, args.train_epochs): if hasattr(dl.sampler, "set_epoch"): dl.sampler.set_epoch(ep) if ba: # PolyLR (B): lr = 1e-2 * (1 - ep/total)^0.9 on the GLOBAL epoch index lr0 = 1e-2 * (1 - ep / args.train_epochs) ** 0.9 for pg in opt.param_groups: pg["lr"] = lr0 else: lr0 = opt.param_groups[0]["lr"] loss_sum = n = 0.0 t_ep = time.time() for batch in dl: img = batch["pimg"].to(dev, non_blocking=True) lab = batch["plabel"].squeeze(1).to(dev, non_blocking=True) wts = batch["w"].to(dev, non_blocking=True) with torch.autocast("cuda", dtype=torch.bfloat16): logits = model(img) if ba: l2 = TF.interpolate(lab.unsqueeze(1).float(), scale_factor=0.5, mode="nearest").squeeze(1).long() l4 = TF.interpolate(lab.unsqueeze(1).float(), scale_factor=0.25, mode="nearest").squeeze(1).long() loss = deep_sup_loss(logits, [lab, l2, l4], b_dc_and_ce) else: loss = loss_f(logits, lab) if (args.c_nn_loss or args.b_loss_only) \ else loss_f(logits, lab, wts) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 12.0 if ba else 5.0) opt.step() if sched is not None: sched.step() loss_sum += float(loss.detach()) n += 1 res.setdefault("per_epoch_loss", []).append(round(loss_sum / max(n, 1), 4)) res.setdefault("per_epoch_time", []).append(round(time.time() - t_ep, 1)) res.setdefault("per_epoch_lr", []).append(round(lr0, 8)) line = (f"[C] epoch {ep + 1}/{args.train_epochs} loss={loss_sum / max(n, 1):.4f} " f"t={time.time() - t_ep:.0f}s lr={lr0:.2e}") if ba and val_dl is not None: vd = val_pseudo_dice(model, val_dl, dev) ema = vd if ema is None else 0.9 * ema + 0.1 * vd res.setdefault("per_epoch_val_dice", []).append(round(vd, 4)) res.setdefault("per_epoch_val_ema", []).append(round(ema, 4)) if best_ema is None or ema > best_ema: best_ema = ema torch.save({"model": model.state_dict(), "epoch": ep + 1, "ema": round(ema, 4)}, best_ema_ckpt) line += f" val={vd:.4f} ema={ema:.4f} best_ema={best_ema:.4f}" print(line, flush=True) if (ep + 1) in eval_at: print(f"[C] mid-eval @ epoch {ep + 1}", flush=True) run_eval(model, f"epoch_{ep + 1}") model.train() if (ep + 1) in eval_at or (ep + 1) % 10 == 0: _dump_partial(res, args) dt = time.time() - t0 res["train_s"] = round(dt, 1) res["train_s_per_epoch"] = round(dt / max(args.train_epochs - start_epoch, 1), 1) res["train_vram_mb"] = peak_mb(dev) res["steps_per_epoch"] = steps if ba: res["best_ema"] = round(best_ema, 4) if best_ema is not None else None res["best_ema_ckpt"] = best_ema_ckpt torch.save({"model": model.state_dict(), "epoch": args.train_epochs, "optimizer": opt.state_dict()}, ckpt) else: torch.save({"model": model.state_dict(), "epoch": args.train_epochs, "optimizer": opt.state_dict(), "scheduler": sched.state_dict()}, ckpt) del dl, opt if sched is not None: del sched res["ckpt"] = res.get("ckpt", ckpt) res["params"] = int(sum(p.numel() for p in model.parameters())) res["params_M"] = round(res["params"] / 1e6, 3) res["weight_mb"] = round(os.path.getsize(res["ckpt"]) / 1e6, 3) model.eval() warm = test_rows[0] predict_probs(model, read_nii_arr(warm["pimg"]).astype("float32"), dev, args.patch, 0.5, tta=tta, sw_batch=8, zscore=ba) torch.cuda.synchronize(dev) reset_peak(dev) dices, times = run_eval(model, "final", timed=True) res["infer_s_per_vol"] = round(float(np.mean(times)), 2) if times else None res["infer_s_per_vol_std"] = round(float(np.std(times)), 2) if times else None res["infer_vram_mb"] = peak_mb(dev) res["dice"] = round(float(np.mean(dices)), 4) if dices else 0.0 res["dice_per_row"] = [{"key": test_rows[args.eval_warmup + i]["key"], "dice": round(dices[i], 4), "s": round(times[i], 2)} for i in range(len(dices))] del model free_gpu() return res # ---------------- Pipeline B (nnU-Net) ---------------- def bench_B(args, train_rows, test_rows, val_rows=None): import nnu_common as N import SimpleITK as sitk if args.skip_train: N.DS_ID = args.nnu_dsid or 210 # study model by default else: N.DS_ID = args.nnu_dsid or 221 # isolated dataset; the study's 210 is never touched N.DS_NAME = "Dataset210_NTUH_T1C_PL" if (N.DS_ID == 210 and args.skip_train) \ else f"Dataset{N.DS_ID:03d}_T1C_BENCH" dev = torch.device(f"cuda:{args.gpu}") torch.cuda.set_device(dev) tta = not args.no_tta res = {"pipeline": "B", "backbone": "nnU-Net 3d_fullres", "tta": tta, "nnu_dataset": N.DS_NAME, "n_train": None, "n_eval": len(test_rows)} te = test_rows if not args.skip_train: tr = train_rows[:args.nnu_cases] nat = native_rows_for([r["key"] for r in tr]) res["n_train"] = len(nat) t0 = time.time() N.make_raw_dataset(nat) res["raw_build_s"] = round(time.time() - t0, 1) if not args.nnu_skip_prep: t0 = time.time() N.run(N.plan_preprocess_cmd(8), os.path.join(d("logs/nnu"), "bench_plan.log"), env=N.nnu_env(), retries=5) # CIFS /mnt/b4 drops during the ~3h preprocess; extra attempts are cheap res["plan_preprocess_s"] = round(time.time() - t0, 1) else: res["plan_preprocess_s"] = 0.0 all_keys = N.dataset_case_keys() val_keys = all_keys[-max(1, len(all_keys) // 10):] N.write_splits(all_keys, val_keys) env = N.nnu_env(epoch=args.train_epochs, lr=1e-2) env["CUDA_VISIBLE_DEVICES"] = str(args.gpu) bcmd = N.train_cmd(1) if args.resume: # nnUNet --c: continue fold_0 from checkpoint_latest.pth (optimizer + # lr-schedule state are restored natively); trains up to PL_EPOCHS # (= --train-epochs, the total budget) bcmd.append("--c") res["resumed_from_checkpoint_latest"] = True try: _st = torch.load(os.path.join(N.fold_dir(), "checkpoint_latest.pth"), map_location="cpu", weights_only=False) res["resumed_from_epoch"] = int(_st.get("current_epoch", _st.get("epoch", 0))) except Exception: pass if not args.nnu_skip_prep: print("[B] --resume without --nnu-skip-prep: plan_and_preprocess " "still runs with --clean (the trained checkpoint is kept)", flush=True) poller = SmiPoller(args.gpu) poller.start() t0 = time.time() N.run(bcmd, os.path.join(d("logs/nnu"), "bench_train.log"), env=env, retries=0) res["train_s"] = round(time.time() - t0, 1) n_trained = max(args.train_epochs - int(res.get("resumed_from_epoch", 0)), 1) res["train_s_per_epoch"] = round(res["train_s"] / n_trained, 1) res["train_vram_mb"] = round(poller.max_mib * 1.048576, 1) if poller.max_mib else None poller.stop() res["train_vram_note"] = "nvidia-smi GPU memory.used peak during nnUNetv2_train sub-process" else: if not os.path.exists(N.best_ckpt()): raise FileNotFoundError( f"--skip-train: {N.best_ckpt()} not found (study nnU-Net results missing)") ckpt = os.path.join(N.fold_dir(), args.nnu_ckpt) res["ckpt"] = ckpt res["weight_mb"] = round(os.path.getsize(ckpt) / 1e6, 1) if os.path.exists(ckpt) else None # in-process inference via nnUNetPredictor os.environ.update(N.nnu_env()) # nnUNet_* roots + nnUNet_extTrainer (NTUHLPLTrainer discovery) from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor predictor = nnUNetPredictor(tile_step_size=0.5, use_gaussian=True, use_mirroring=tta, perform_everything_on_device=True, device=dev, verbose=False, allow_tqdm=False) predictor.initialize_from_trained_model_folder(N.model_folder(), use_folds=[0], checkpoint_name=args.nnu_ckpt) net = predictor.network if hasattr(net, "module"): net = net.module res["params"] = int(sum(p.numel() for p in net.parameters())) res["params_M"] = round(res["params"] / 1e6, 3) te_keys = [r["key"] for r in te] te_native = native_rows_for(te_keys) pred_base = os.environ.get("BENCH_B_DIR", os.path.join(d("runs_bench"), "B")) in_dir = os.path.join(pred_base, "pred_in") out_dir = os.path.join(pred_base, "pred_out") shutil.rmtree(in_dir, ignore_errors=True) shutil.rmtree(out_dir, ignore_errors=True) os.makedirs(in_dir) os.makedirs(out_dir) for r in te_native: os.symlink(os.path.abspath(r["img"]), os.path.join(in_dir, r["key"] + "_0000.nii.gz")) def run_cases(cs): # list-of-lists: one channel-file list per case imgs = [[os.path.join(in_dir, c["key"] + "_0000.nii.gz")] for c in cs] predictor.predict_from_files(imgs, out_dir, save_probabilities=True, overwrite=True, num_processes_preprocessing=2, num_processes_segmentation_export=2, num_parts=1, part_id=0) warm, timed = te_native[:args.eval_warmup], te_native[args.eval_warmup:] if warm: run_cases(warm) torch.cuda.synchronize(dev) reset_peak(dev) t0 = time.time() if timed: run_cases(timed) res["infer_s_per_vol"] = round((time.time() - t0) / len(timed), 2) if timed else None res["infer_vram_mb"] = peak_mb(dev) if timed else None dices, dice_hard, per = [], [], [] for r in timed: key = r["key"] npz_f = os.path.join(out_dir, key + ".npz") seg_f = os.path.join(out_dir, key + ".nii.gz") if not os.path.exists(npz_f): per.append({"key": key, "dice": None, "note": "no npz"}) continue probs = np.load(npz_f, allow_pickle=False)["probabilities"] if probs.shape[0] != 2: per.append({"key": key, "dice": None, "note": f"bad channels {probs.shape[0]}"}) continue pt = probs[1] img_itk = sitk.ReadImage(r["img"]) lab = load_native_label(r["label"], img_itk) if pt.shape != lab.shape: per.append({"key": key, "dice": None, "note": f"prob {pt.shape} vs label {lab.shape}"}) continue dices.append(dice_np(pt, lab)) if os.path.exists(seg_f): hard = read_nii_arr(seg_f) if hard.shape == lab.shape: dice_hard.append(dice_np((hard > 0).astype("float32"), lab)) per.append({"key": key, "dice": round(dices[-1], 4)}) res["dice_per_row"] = per res["dice"] = round(float(np.mean(dices)), 4) if dices else 0.0 if dice_hard: res["dice_hard"] = round(float(np.mean(dice_hard)), 4) shutil.rmtree(in_dir, ignore_errors=True) del predictor free_gpu() return res # ---------------- report ---------------- def print_table(results): print() hdr = (f"{'pipeline':46s} {'dice':>6s} {'paramsM':>8s} {'wtMB':>7s} " f"{'s/epoch':>8s} {'trVRAM':>8s} {'s/vol':>7s} {'inVRAM':>8s}") print(hdr) print("-" * len(hdr)) for r in results: if "error" in r: print(f"{r['pipeline']:46s} ERROR {r['error']}") continue extra = [] if r["pipeline"] == "B": extra.append(f"prep={r.get('plan_preprocess_s')}s") line = (f"{r['pipeline']+' '+r['backbone'][:40]:46s} " f"{r['dice']:6.4f} {r['params_M']:8.2f} {r.get('weight_mb', 0):7.2f} " f"{r.get('train_s_per_epoch', float('nan')):8.1f} " f"{(r.get('train_vram_mb') or float('nan')):8.0f} " f"{(r.get('infer_s_per_vol') or float('nan')):7.2f} " f"{(r.get('infer_vram_mb') or float('nan')):8.0f}") if extra: line += " " + " ".join(extra) print(line) print() def _dump_partial(res, args): """Write the in-progress res dict to .partial.json (crash protection: mid-training evals + per-epoch loss survive a failed run; atomic replace).""" try: out_f = args.out or os.path.join(ROOT, "results", "benchmark_pipelines", "benchmark_pipelines.json") payload = {"timestamp": datetime.now().isoformat(timespec="seconds"), "partial": True, "settings": {k: v for k, v in vars(args).items()}, "results": [res]} tmp = out_f + ".partial.json.tmp" with open(tmp, "w") as f: json.dump(payload, f, indent=1) os.replace(tmp, out_f + ".partial.json") except Exception as e: # noqa print(f"partial dump failed: {e!r}", flush=True) def _reconstructed_lr(steps_per_epoch, epochs_total, start_epoch, base_lr=3e-4, eta_frac=0.05): """LR at the start of each epoch for the A/C schedule (3e-4, 300-step LinearLR warmup from 0.1x, then CosineAnnealing to base_lr*eta_frac over the TOTAL run's steps) — used to draw the lr panel for runs that predate per-epoch LR logging.""" total = steps_per_epoch * epochs_total warmup = min(300, max(10, total // 10)) eta = base_lr * eta_frac out = [] for ep in range(start_epoch + 1, epochs_total + 1): s = (ep - 1) * steps_per_epoch if s < warmup: lr = base_lr * (0.1 + 0.9 * s / max(warmup, 1)) else: t = min(s - warmup, max(total - warmup, 1)) lr = eta + 0.5 * (base_lr - eta) * (1 + math.cos(math.pi * t / max(total - warmup, 1))) out.append(lr) return out def plot_progress(r, png): """nnU-Net-style progress figure (like fold_0/progress.png): stacked panels 1) train loss (left) + test dice (right axis, from res['evals']), 2) epoch duration (when per-epoch times are logged), 3) learning rate (logged, or reconstructed for the A/C schedule).""" try: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt try: plt.style.use("seaborn-v0_8-darkgrid") except OSError: plt.style.use("darkgrid") except Exception as e: # noqa print("progress plot failed:", repr(e)) return loss = r.get("per_epoch_loss") or [] start = int(r.get("resumed_from_epoch", 0)) + 1 epochs = list(range(start, start + len(loss))) ev = r.get("evals") or {} mid = sorted((int(t.split("_")[1]), d["dice"]) for t, d in ev.items() if t.startswith("epoch_")) final_ep = int(r.get("epochs_total", epochs[-1] if epochs else (mid[-1][0] if mid else 0))) dice_pts = mid + ([(final_ep, ev["final"]["dice"])] if "final" in ev else []) dur = r.get("per_epoch_time") or [] if len(dur) != len(epochs): dur = [] lr = r.get("per_epoch_lr") or [] if len(lr) != len(epochs): lr = _reconstructed_lr(int(r.get("steps_per_epoch", 0)), final_ep, start - 1) \ if r.get("steps_per_epoch") and final_ep else [] have = [p for p in (("loss", bool(loss)), ("dur", bool(dur)), ("lr", bool(lr))) if p[1]] if not have: return blue, green = "#1f77b4", "#2ca02c" fig, axes = plt.subplots(len(have), 1, figsize=(9, 4.8 * len(have)), sharex=True) if len(have) == 1: axes = [axes] for ax, (name, _) in zip(axes, have): if name == "loss": ax.plot(epochs, loss, color=blue, lw=1.0, label="loss_tr") ax.set_ylabel("loss") handles, labels = ax.get_legend_handles_labels() if dice_pts: ax2 = ax.twinx() ax2.plot([p[0] for p in dice_pts], [p[1] for p in dice_pts], color=green, lw=1.5, marker="o", ms=4, label="test dice (TTA)") ax2.set_ylabel("test dice") handles += ax2.get_legend_handles_labels()[0] labels += ax2.get_legend_handles_labels()[1] ax.legend(handles, labels, loc="upper left", fontsize=8) elif name == "dur": ax.plot(epochs, dur, color=blue, lw=1.0, label="epoch duration") ax.set_ylabel("time [s]") ax.legend(loc="upper left", fontsize=8) elif name == "lr": ax.plot(epochs, lr, color=blue, lw=1.2, label="learning rate") ax.set_ylabel("learning rate") ax.legend(loc="upper left", fontsize=8) ax.set_xlabel("epoch") fig.suptitle(f"pipeline {r['pipeline']} {r.get('backbone', '')} — {r.get('loss', '')}", fontsize=10) fig.tight_layout() fig.savefig(png, dpi=130) plt.close(fig) print("progress.png:", png) def plot(results, out_png, epochs): try: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt ok = [r for r in results if "error" not in r] if not ok: return names = [f"{r['pipeline']} {r['backbone']}" for r in ok] fig, axes = plt.subplots(2, 3, figsize=(15, 8)) panels = [ ("DICE (test", "dice", "{:.4f}"), ("Train s/epoch", "train_s_per_epoch", " {:.0f}"), ("Infer s/vol (TTA as marked)", "infer_s_per_vol", " {:.1f}"), ("Params (M)", "params_M", " {:.1f}"), ("VRAM train (MB)", "train_vram_mb", " {:.0f}"), ("VRAM infer (MB)", "infer_vram_mb", " {:.0f}"), ] for ax, (title, key, fmt) in zip(axes.ravel(), panels): vals = [r.get(key) for r in ok] vals = [v if v is not None else 0.0 for v in vals] bars = ax.bar(range(len(ok)), vals, color=["#4878CF", "#6ACC65", "#D64F5C"][:len(ok)]) for b, v in zip(bars, vals): ax.annotate(fmt.format(v), (b.get_x() + b.get_width() / 2, b.get_height()), ha="center", va="bottom", fontsize=9) ax.set_title(title + ")", fontsize=10) ax.set_xticks(range(len(ok))) ax.set_xticklabels([n.split(" ")[0] for n in names], fontsize=10) plt.suptitle(f"Pipeline benchmark: {epochs} train epochs, single GPU", fontsize=11) fig.tight_layout() fig.savefig(out_png, dpi=130) plt.close(fig) print("plot:", out_png) except Exception as e: # noqa print("plot failed:", repr(e)) def main(): args = parse_args() if args.std_loss: args.nn_loss = False # Per-pipeline config mode. B-aligned is C's DEFAULT (its record config) # and A's OPT-IN (--b-aligned); A's DEFAULT is A's own defaults (the # 0.5508 @100ep record, see BENCHMARKING.md). --legacy and --b-loss-only # opt out for both; --skip-train evaluates the legacy-config study ckpts. if args.skip_train or args.legacy or args.b_loss_only: args.a_ba = args.c_ba = False else: args.c_ba = True args.a_ba = bool(args.b_aligned) if args.nn_loss is None: # no explicit loss flag: A's default mode uses its own std loss; # --legacy / --skip-train keep the historical approx nn-loss default args.a_nn_loss = args.legacy or args.skip_train args.c_nn_loss = True else: args.a_nn_loss = args.c_nn_loss = bool(args.nn_loss) if args.skip_train: print("[bench] --skip-train: A/C study checkpoints are legacy-config; " "evaluating with the legacy input protocol", flush=True) pipelines = [p.upper() for p in args.pipelines.split()] unknown = set(pipelines) - set("ABC") if unknown: raise SystemExit(f"unknown pipeline(s): {unknown}") out_f = args.out or os.path.join(ROOT, "results", "benchmark_pipelines", "benchmark_pipelines.json") os.makedirs(os.path.dirname(out_f), exist_ok=True) train_rows, test_rows, val_rows = select_rows(args) a_mode = ("b_aligned" if args.a_ba else ("b_loss_only" if args.b_loss_only else ("legacy" if args.legacy or args.a_nn_loss else "a_default"))) c_mode = ("b_aligned" if args.c_ba else ("b_loss_only" if args.b_loss_only else "legacy")) print(f"[bench] gpu={args.gpu} train_rows={len(train_rows)} eval_rows={len(test_rows)} " f"val_rows={len(val_rows)} epochs={args.train_epochs} tta={not args.no_tta} " f"skip_train={args.skip_train} A config={a_mode} C config={c_mode}") runners = {"A": bench_A, "B": bench_B, "C": bench_C} results = [] for p in pipelines: t0 = time.time() print(f"\n[bench] ---- pipeline {p} ----") try: r = runners[p](args, train_rows, test_rows, val_rows) r["wall_s"] = round(time.time() - t0, 1) except Exception as e: # noqa tb = traceback.format_exc() r = {"pipeline": p, "error": repr(e), "traceback": tb, "wall_s": round(time.time() - t0, 1)} print(f"[bench] pipeline {p} FAILED:\n{tb}") results.append(r) free_gpu() payload = { "timestamp": datetime.now().isoformat(timespec="seconds"), "settings": {k: v for k, v in vars(args).items()}, "protocol": ("single GPU; bf16 batch 3 patch 96 for A/C; A default = A-native " "all-defaults (BN + ReLU + per-sample 0.5CE+0.5dice + AdamW 3e-4 " "warmup/cosine + clip 5 + legacy aug + percentile + 50% fg), " "A --b-aligned = B's protocol (instance-norm + deep sup + B's DC+CE, " "SGD 1e-2 PolyLR, B aug, Z-score), --legacy = pre-2026-10-07 config " "(approx nn-loss default); C default = B-aligned; A's TTA blending " "fixed 2026-10-06; B nnU-Net planned batch/patch fp16 SGD 1e-2 " "PolyLR; dice = prob@0.5 vs GT; inference s/vol = mean over timed " "test cases after warmup; train VRAM A/C = torch peak, B = nvidia-smi " "peak; B dataset id isolated at " f"{args.nnu_dsid} in train mode"), "results": results, } with open(out_f, "w") as f: json.dump(payload, f, indent=1) try: os.remove(out_f + ".partial.json") except OSError: pass per_rows = [x for r in results for x in r.get("dice_per_row", [])] if per_rows: with open(out_f.replace(".json", "_per_row.jsonl"), "w") as f: for x in per_rows: f.write(json.dumps(x) + "\n") print_table(results) print("json:", out_f) plot(results, out_f.replace(".json", "_bench.png"), args.train_epochs) # per-pipeline progress.png (train loss + mid/final test dice); B's comes # from nnU-Net's own plot next to its checkpoint for r in results: if "error" in r: continue png = out_f.replace(".json", "") + f"_{r['pipeline']}_progress.png" if r.get("per_epoch_loss"): plot_progress(r, png) elif r.get("ckpt"): src = os.path.join(os.path.dirname(r["ckpt"]), "progress.png") if os.path.exists(src): shutil.copyfile(src, png) print("progress.png:", png) if __name__ == "__main__": main()