feat(benchmarking): add pipeline benchmarking tools

Add a new benchmarking script and documentation to facilitate
performance comparisons between the different implemented pipelines.

- Create `scripts/benchmark_pipelines.py` for automated evaluation
- Add `BENCHMARKING.md` to outline benchmarking procedures and metrics

BENCHMARKING.md

@@ -0,0 +1,151 @@
+# Benchmarking Pipelines A / B / C
+
+`scripts/benchmark_pipelines.py` compares the three segmentation pipelines on
+**one GPU** with a unified protocol:
+
+| pipeline | backbone | data protocol |
+|---|---|---|
+| A | in-house `Unet3D(base=16, depth=4)` | processed 1 mm head-cropped percentile-normalized volumes (`pimg`/`plabel`) |
+| B | nnU-Net v2 `3d_fullres` | native source volumes (`img`/`label`), per Pipeline B's design |
+| C | MONAI `UNet(16→128)` | processed 1 mm volumes (`pimg`/`plabel`) |
+
+Each pipeline is benchmarked with **its own standard settings** (batch, patch,
+augmentation, TTA, optimizer), so the results reflect the pipelines as used in
+the study.
+
+## Metrics
+
+- **DICE** — tumor probability thresholded at 0.5 vs ground truth on the
+  held-out patient-level test split (`split_test.jsonl`), the primary-metric
+  convention of all three pipelines. B additionally reports the **hard-seg**
+  (argmax) Dice it exports natively.
+- **Training** — wall-clock total and s/epoch for `--train-epochs` on the
+  first `--max-rows` train rows (single GPU). For B, the one-time raw-dataset
+  build and `plan_and_preprocess` wall time are reported separately
+  (`raw_build_s`, `plan_preprocess_s`).
+- **Inference** — mean **s/volume** on `--max-eval-rows` held-out test volumes
+  (first `--eval-warmup` cases untimed), each pipeline's standard
+  sliding-window + TTA protocol (`--no-tta` disables TTA everywhere).
+- **Model size** — parameter count, checkpoint weight file size (MB), and peak
+  **VRAM** (MB) for the training run and for the inference run.
+  (A/C training VRAM via `torch.cuda.max_memory_allocated`; B via nvidia-smi
+  polling around the `nnUNetv2_train` subprocess.)
+
+## Prerequisites
+
+- The `longitudinal` conda env and the GPU(s) it sees:
+  ```bash
+  source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
+  ```
+- Data for Pipelines A and C (and split membership for all three):
+  `data/proc/*` + `data/manifests/split_{train,val,test}.jsonl`, i.e.
+  `scripts/preprocess.py` and `scripts/05_build_splits.py` have been run.
+- Data for Pipeline B (train/eval mode): the labeled source manifests
+  `data/manifests/{ntuh,m6_labeled}.jsonl` with reachable native NIfTIs
+  (the benchmark joins native `img`/`label` paths by case key).
+- `--skip-train` mode additionally needs the study checkpoints:
+  `runs/round0/best.pt` (A), `runs_monai/round0/best.pt` (C), and the study
+  nnU-Net results tree `nnu/results/Dataset210_NTUH_T1C_PL/...` (B).
+
+## Quick start
+
+```bash
+cd /mnt/b4/xfr/git26/longitudinal
+conda activate longitudinal
+
+# small/fast check (what the validation run used): ~minutes per pipeline
+python scripts/benchmark_pipelines.py \
+    --train-epochs 1 --max-rows 24 --nnu-cases 32 --max-eval-rows 4
+
+# default full benchmark: train 2 epochs on 120 rows, eval on 12 test volumes
+python scripts/benchmark_pipelines.py
+```
+
+Outputs (repo root):
+
+```
+results/benchmark_pipelines.json        # settings + all metrics, per pipeline
+results/benchmark_pipelines_per_row.jsonl
+results/benchmark_pipelines_bench.png   # 6-panel bar chart
+```
+
+plus the artifacts it created (safe to delete): `runs_bench/{A,C}/best.pt`,
+`nnu/{raw,preprocessed,results}/Dataset221_T1C_BENCH`,
+`logs/nnu/bench_{plan,train}.log`.
+
+## All options
+
+| flag | default | meaning |
+|---|---|---|
+| `--pipelines` | `A B C` | space-separated subset, e.g. `"A C"` |
+| `--gpu` | `0` | physical GPU index; all pipelines run serially on this one GPU |
+| `--train-epochs` | `2` | training budget for A and C (and B via `NNU_PL_EPOCHS`) |
+| `--max-rows` | `120` | train rows for A/C (first N of `split_train.jsonl` that have labels) |
+| `--nnu-cases` | `120` | train cases for B (first N labeled native cases of the same split) |
+| `--max-eval-rows` | `12` | held-out test volumes for inference + DICE |
+| `--eval-warmup` | `1` | untimed warmup cases before inference timing |
+| `--workers` | `4` | data-loader workers (A/C) |
+| `--batch` | `3` | per-GPU batch (A/C) |
+| `--patch` | `96` | patch/sliding-window size (A/C) |
+| `--no-tta` | off | disable each pipeline's standard test-time augmentation |
+| `--skip-train` | off | evaluate existing study checkpoints instead of training |
+| `--nnu-dsid` | `221` train / `210` skip-train | nnU-Net dataset id: in training mode an isolated `Dataset{id}_T1C_BENCH` (the study's 210 is never touched); in `--skip-train` mode the study's 210 by default, or a previously benchmarked id to re-evaluate it |
+| `--nnu-skip-prep` | off | B (train mode): skip `plan_and_preprocess` and reuse the existing raw + preprocessed tree for `--nnu-dsid` (e.g. retime one more training epoch on an already prepared dataset) |
+| `--out` | `results/benchmark_pipelines.json` | output JSON path (the `_per_row.jsonl` / `_bench.png` are derived from it) |
+
+## Examples
+
+```bash
+# longer training budget, bigger eval set (quality-focused run)
+python scripts/benchmark_pipelines.py --train-epochs 5 --max-rows 240 --max-eval-rows 20
+
+# timing only, no TTA
+python scripts/benchmark_pipelines.py --no-tta --max-eval-rows 20
+
+# re-evaluate the study's trained checkpoints (no training)
+python scripts/benchmark_pipelines.py --skip-train --max-eval-rows 24
+
+# benchmark B only; keep the trained Dataset221 model for a second inference-only pass
+python scripts/benchmark_pipelines.py --pipelines B
+python scripts/benchmark_pipelines.py --pipelines B --skip-train --nnu-dsid 221
+
+# use GPU 1 and write to a custom location
+python scripts/benchmark_pipelines.py --gpu 1 --out results/bench_gpu1.json
+```
+
+## Interpreting the results
+
+- **"1 epoch" is pipeline-native.** A/C: one pass over the N train volumes
+  (batch 3, 96³ random patches, AdamW 3e-4 + warmup/cosine, bf16).
+  B: one pass over the N train cases at nnU-Net's planned batch and patch size
+  (typically 2 and ~256³, AdamW 1e-2 PolyLR, fp16). Epoch times are therefore
+  comparable as *wall-clock cost per pass over the training set in each
+  pipeline's standard configuration*, not per identical number of patches.
+- **DICE reflects the training budget.** Short benchmarks measure the
+  machinery, not the final quality of the pipelines. Compare DICE across
+  pipelines only at the same `--train-epochs`/`--max-rows` budget, or use
+  `--skip-train` on fully trained study checkpoints.
+- **Inference times are single-GPU, standard-protocol numbers**
+  (A: 96³ step 48 + 4-flip TTA; B: nnU-Net gaussian sliding window +
+  mirroring, native volumes; C: MONAI gaussian 50% + 4-flip TTA).
+  B is the slowest per volume on this data set because it runs on uncropped
+  native volumes with full mirroring TTA.
+- **B's one-time cost.** `raw_build_s` + `plan_preprocess_s` are paid once per
+  (dataset, round) — in the iterative study every round repays the
+  plan/preprocess on the growing dataset; the benchmark reports it for
+  context only.
+- **VRAM notes.** Training VRAM includes optimizer states (AdamW) and the
+  largest batch; B's number is the nvidia-smi peak of the whole
+  `nnUNetv2_train` process tree on that GPU. Inference VRAM is per
+  pipeline as measured in-process (B: torch peak over the timed cases).
+
+## Troubleshooting
+
+| symptom | fix |
+|---|---|
+| `split_train.jsonl has no labeled processed rows` | run `scripts/preprocess.py` over the manifests and `scripts/05_build_splits.py` first |
+| `B: key ... not in ntuh/m6_labeled manifests` | build the source manifests (`scripts/01_build_ntuh_manifest.py`, `scripts/02_build_m6_dataset.py`) and confirm the native NIfTIs are reachable |
+| `--skip-train: .../best.pt not found` | the study checkpoint for that pipeline doesn't exist yet — train a round first, or drop `--skip-train` |
+| B `--skip-train` fails on the model folder | the study's `nnu/results/Dataset210_NTUH_T1C_PL/NTUHLPLTrainer__nnUNetPlans__3d_fullres/fold_0/` must exist (produced by `scripts_nnu/03_nnu_train.py`) |
+| OOM on B | B trains at the planned 256³-class patches and needs ~10–20 GB; free other GPU users or run with a smaller `--nnu-cases` set (patch size derives from the median case size) |
+| inflated/inconsistent timings | do not run two benchmarks on the same `--gpu` concurrently; note that data-loader (NFS) latency is included in training and inference wall times |

scripts/benchmark_pipelines.py

@@ -0,0 +1,641 @@
+"""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). "1 epoch" = one pass over the training
+  set under each pipeline's native settings: A/C batch 3, 96^3 random
+  patches, AdamW 3e-4 + warmup/cosine; B nnU-Net planned batch (2), planned
+  patch size, AdamW 1e-2 PolyLR. 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` 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 --no-tta                 # no test-time augmentation
+  python scripts/benchmark_pipelines.py --skip-train             # eval existing study checkpoints
+Outputs: results/benchmark_pipelines.json (+ _per_row.jsonl, _bench.png)
+"""
+import argparse
+import gc
+import json
+import os
+import shutil
+import subprocess
+import sys
+import threading
+import time
+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 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=2, help="training budget for A and C (and B via NNU_PL_EPOCHS)")
+    ap.add_argument("--max-rows", type=int, default=120, help="train rows for A/C (first N of split_train)")
+    ap.add_argument("--nnu-cases", type=int, default=120, help="train cases for B (first N labeled native cases)")
+    ap.add_argument("--max-eval-rows", type=int, default=12, help="held-out test volumes for inference + DICE")
+    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("--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("--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("--out", default=None, help="output json path (default results/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(ROOT, "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(ROOT, "data", "manifests")
+    train = load_jsonl(os.path.join(man, "split_train.jsonl"))
+    test = load_jsonl(os.path.join(man, "split_test.jsonl"))
+    train = [r for r in train if r.get("plabel")][:args.max_rows]
+    test = [r for r in test if r.get("plabel")][:args.max_eval_rows]
+    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")
+    return train, test
+
+
+# ---------------- 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):
+    from src.dataset import make_dataloader
+    from src.losses import per_sample_loss, one_hot
+    from src import training as A
+
+    dev = torch.device(f"cuda:{args.gpu}")
+    torch.cuda.set_device(dev)
+    tta = not args.no_tta
+    res = {"pipeline": "A", "backbone": "Unet3D(base=16, depth=4)", "tta": tta,
+           "n_train": len(train_rows), "n_eval": len(test_rows)}
+
+    ckpt_dir = d("runs_bench/A")
+    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:
+        model = A.build_model(base=16, device="cuda")
+        patch = (args.patch, args.patch, args.patch)
+        _, dl = make_dataloader(train_rows, patch, args.batch, True, num_workers=args.workers,
+                                seed=0, persistent=False)
+        steps = max(len(dl), 1)
+        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])
+        model.train()
+        reset_peak(dev)
+        t0 = time.time()
+        for ep in range(args.train_epochs):
+            if hasattr(dl.sampler, "set_epoch"):
+                dl.sampler.set_epoch(ep)
+            loss_sum = n = 0.0
+            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)
+                losses = per_sample_loss(logits.float(), one_hot(lab, 2))
+                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)
+                opt.zero_grad(set_to_none=True)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
+                opt.step()
+                sched.step()
+                loss_sum += float(losses.detach().mean())
+                n += 1
+            res.setdefault("per_epoch_loss", []).append(round(loss_sum / max(n, 1), 4))
+        dt = time.time() - t0
+        res["train_s"] = round(dt, 1)
+        res["train_s_per_epoch"] = round(dt / max(args.train_epochs, 1), 1)
+        res["train_vram_mb"] = peak_mb(dev)
+        res["steps_per_epoch"] = steps
+        ckpt = os.path.join(ckpt_dir, "best.pt")
+        torch.save({"model": model.state_dict(), "epoch": args.train_epochs}, ckpt)
+        del dl, opt, sched
+    res["ckpt"] = res.get("ckpt", os.path.join(ckpt_dir, "best.pt"))
+    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)
+    torch.cuda.synchronize(dev)
+    reset_peak(dev)
+    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()
+            probs = A.sliding_window_probs(model, vol, dev, args.patch, step, tta)
+            torch.cuda.synchronize(dev)
+            times.append(time.time() - t0)
+            # probs are on the padded grid; crop back to the volume/label shape
+            osh = vol.shape
+            probs = probs[:osh[0], :osh[1], :osh[2]]
+            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))
+    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 C (MONAI) ----------------
+
+def bench_C(args, train_rows, test_rows):
+    sys.path.insert(0, os.path.join(ROOT, "scripts_monai"))
+    from monai_common import (build_model, WeightedDiceCELoss, make_dataloader,
+                              predict_probs, prob_dice, load_model)
+
+    dev = torch.device(f"cuda:{args.gpu}")
+    torch.cuda.set_device(dev)
+    tta = not args.no_tta
+    res = {"pipeline": "C", "backbone": "monai UNet(16->128)", "tta": tta,
+           "n_train": len(train_rows), "n_eval": len(test_rows)}
+
+    ckpt_dir = d("runs_bench/C")
+    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)
+        dl = make_dataloader(train_rows, win=args.patch, batch=args.batch, workers=args.workers)
+        steps = max(len(dl), 1)
+        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])
+        loss_f = WeightedDiceCELoss()
+        model.train()
+        reset_peak(dev)
+        t0 = time.time()
+        for ep in range(args.train_epochs):
+            if hasattr(dl.sampler, "set_epoch"):
+                dl.sampler.set_epoch(ep)
+            loss_sum = n = 0.0
+            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)
+                loss = loss_f(logits, lab, wts)
+                opt.zero_grad(set_to_none=True)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
+                opt.step()
+                sched.step()
+                loss_sum += float(loss.detach())
+                n += 1
+            res.setdefault("per_epoch_loss", []).append(round(loss_sum / max(n, 1), 4))
+        dt = time.time() - t0
+        res["train_s"] = round(dt, 1)
+        res["train_s_per_epoch"] = round(dt / max(args.train_epochs, 1), 1)
+        res["train_vram_mb"] = peak_mb(dev)
+        res["steps_per_epoch"] = steps
+        ckpt = os.path.join(ckpt_dir, "best.pt")
+        torch.save({"model": model.state_dict(), "epoch": args.train_epochs}, ckpt)
+        del dl, opt, sched, loss_f
+    res["ckpt"] = res.get("ckpt", os.path.join(ckpt_dir, "best.pt"))
+    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)
+    torch.cuda.synchronize(dev)
+    reset_peak(dev)
+    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(model, vol, dev, args.patch, 0.5, tta=tta, sw_batch=8)
+            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))
+    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):
+    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())
+            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)
+        poller = SmiPoller(args.gpu)
+        poller.start()
+        t0 = time.time()
+        N.run(N.train_cmd(1), os.path.join(d("logs/nnu"), "bench_train.log"), env=env, retries=0)
+        res["train_s"] = round(time.time() - t0, 1)
+        res["train_s_per_epoch"] = round(res["train_s"] / max(args.train_epochs, 1), 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 = N.best_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="checkpoint_best.pth")
+    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)
+    in_dir = d("runs_bench/B/pred_in")
+    out_dir = d("runs_bench/B/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 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()
+    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.json")
+    os.makedirs(os.path.dirname(out_f), exist_ok=True)
+    train_rows, test_rows = select_rows(args)
+
+    print(f"[bench] gpu={args.gpu} train_rows={len(train_rows)} eval_rows={len(test_rows)} "
+          f"epochs={args.train_epochs} tta={not args.no_tta} skip_train={args.skip_train}")
+
+    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)
+            r["wall_s"] = round(time.time() - t0, 1)
+        except Exception as e:  # noqa
+            r = {"pipeline": p, "error": repr(e), "wall_s": round(time.time() - t0, 1)}
+            print(f"[bench] pipeline {p} FAILED: {e!r}")
+        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; A/C bf16 batch 3 patch 96 AdamW 3e-4; B nnU-Net planned "
+                     "batch/patch fp16 AdamW 1e-2; 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)
+    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)
+
+
+if __name__ == "__main__":
+    main()