diff --git a/BENCHMARKING.md b/BENCHMARKING.md new file mode 100644 index 0000000..008d052 --- /dev/null +++ b/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 | \ No newline at end of file diff --git a/scripts/benchmark_pipelines.py b/scripts/benchmark_pipelines.py new file mode 100644 index 0000000..f3e1ba7 --- /dev/null +++ b/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() \ No newline at end of file