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