Furen Xiao
8fe1818c7d
- Updated regex patterns for improved matching of series names and tags. - Added functionality to reject non-head series based on study and series descriptions. - Implemented max voxel spacing check to filter out series with excessive spacing. - Enhanced the reconstruction script to handle dynamic-frame series exclusions and artifact pruning. - Modified output paths for reconstructed NIfTI files and added QA screenshot generation. - Improved argument parsing in benchmark and pseudo-labeling scripts for better flexibility. - Introduced a new script for generating QA screenshots from reconstructed volumes.
648 lines
No EOL
28 KiB
Python
648 lines
No EOL
28 KiB
Python
"""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=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("--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 = [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")]
|
|
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]
|
|
|
|
|
|
# ---------------- 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() |