diff --git a/BENCHMARKING.md b/BENCHMARKING.md index 1c86b9a..4889866 100644 --- a/BENCHMARKING.md +++ b/BENCHMARKING.md @@ -1,152 +1,96 @@ -# Benchmarking Pipelines A / B / C +# Pipeline benchmark (A / B / C) -`scripts/benchmark_pipelines.py` compares the three segmentation pipelines on -**one GPU** with a unified protocol: +Benchmark of the three segmentation pipelines — A in-house Unet3D, B nnU-Net v2 +(NTUHLPLTrainer), C MONAI UNet — driven by `scripts/benchmark_pipelines.py` +(DICE, training time, inference time, params, VRAM). Protocol details live in +the script docstring; this file records what has been run and what the machine +currently supports. -| 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`) | +## Data chain status (2026-10-04) -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. +The M6 labeled set was fully rebuilt on 2026-10-04 (see +`scripts/02_build_m6_dataset.py`): T1c fixed, best-of-ANTs +(`ants_rigid` / `ants_affine` / `ants_syn`, winner by Mattes MI), TV warped +onto the T1c grid with SMD+linear. After the repair pass: -## Metrics +- `data/manifests/m6_labeled.jsonl`: **1312** volumes (was 1094 pre-repair), + `m6_unlabeled.jsonl`: **229** +- all 1541 m6 processed volumes (`data/proc/{,_label}.nii.gz`, 1mm + cropped/normalized) regenerated, 0 errors +- patient-level splits rebuilt (`scripts/05_build_splits.py`, seed 0): + **4200** labeled rows (ntuh + m6) → train 3570 / val 211 / test 419 + volumes, 3893 subjects; unlabeled pool 13280 +- tumor volume stats: median 3116 mm3, p95 24.8k, max 3.2M, zero_frac 0.12% -- **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.) +## Smoke run (2026-10-04) -## Prerequisites - -- The `longitudinal-313` conda env and the GPU(s) it sees: - ```bash - source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 - ``` -- 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-313 - -# 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 100 epochs on 80% of train cases (~1986 rows), -# eval on 10% of test volumes (~30) -python scripts/benchmark_pipelines.py -``` - -Outputs (repo root): +Small/fast end-to-end check that all three pipelines run on the rebuilt data +(1 epoch, 24 train rows, 3 eval rows, TTA on, GPU 0): ``` -results/benchmark_pipelines.json # settings + all metrics, per pipeline -results/benchmark_pipelines_per_row.jsonl -results/benchmark_pipelines_bench.png # 6-panel bar chart +python scripts/benchmark_pipelines.py --gpu 0 --max-rows 24 --nnu-cases 24 \ + --max-eval-rows 3 --train-epochs 1 --out results/benchmark_pipelines_smoke.json ``` -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`. +| pipeline | wall | s/epoch | tr VRAM | s/vol (TTA) | in VRAM | params | weights | dice (1 epoch) | +|---|---|---|---|---|---|---|---|---| +| A Unet3D(16, depth 4) | 17 s | 2.5 (8 steps) | 2.0 GB | 4.12 | 0.5 GB | 1.46 M | 5.9 MB | 0.0000 | +| B nnU-Net 3d_fullres | 285 s (incl. 83 s plan&preprocess, 4 s raw build) | 73.3 | 9.8 GB | 34.73 | 1.8 GB | 30.8 M | 246 MB | 0.0000 | +| C MONAI UNet(16→128) | 5 s | 2.5 (8 steps) | 1.7 GB | 0.33 | 0.5 GB | 1.19 M | 4.8 MB | 0.0005 | -## All options +DICE ≈ 0 is expected (models trained for a single epoch on 24 volumes); the +point of the run was validating the full chain — new m6 native/processed +manifests, split membership, B's `ntuh + m6_labeled` native join, nnU-Net +isolated dataset 221 build, inference + dice scoring. All three passed. +Eval rows for the smoke run were the first three `split_test` volumes (all +ntuh). Artifacts: `results/benchmark_pipelines_smoke.json`, +`_per_row.jsonl`, `_bench.png`. -| 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` | `100` | training budget for A and C (and B via `NNU_PL_EPOCHS`) | -| `--max-rows` | `80%` of `split_train` | train rows for A/C (first N of `split_train.jsonl` that have labels; default 80% of the split, ~1986 rows) | -| `--nnu-cases` | `80%` of `split_train` | train cases for B (first N labeled native cases of the same split; default same as `--max-rows`) | -| `--max-eval-rows` | `10%` of `split_test` | held-out test volumes for inference + DICE (default 10% of the split, ~30 rows) | -| `--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) | +### Environment fixes made during the check -## Examples +- `monai` was missing from `longitudinal-313` (pipeline C failed with + ModuleNotFoundError) — installed `pip install monai==1.6.1 --no-deps` + (torch 2.14+cu126 untouched); `import monai` verified. -```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 +## Resuming training (`--resume`) -# timing only, no TTA -python scripts/benchmark_pipelines.py --no-tta --max-eval-rows 20 +All three pipelines can continue from their latest benchmark checkpoint: -# 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 +``` +python scripts/benchmark_pipelines.py --pipelines "A B C" --resume \ + --train-epochs 200 --nnu-skip-prep ``` -## Interpreting the results +- `--train-epochs` is the **total budget** (not epochs added). Each pipeline + trains from where its checkpoint left off up to that total. +- **A / C**: load `runs_bench/{A,C}/best.pt` (model + optimizer + scheduler + state — checkpoints written before 2026-10-04 carry model + epoch only and + fall back to fresh optimizer/scheduler) and run the remaining epochs; the + run is a no-op train (eval only) when the checkpoint already reached the + total. +- **B**: appends nnU-Net's `--c` flag (continue `fold_0` from + `checkpoint_latest.pth`, native optimizer + LR-schedule restore) and + `NNUNet_pl_EPOCHS = --train-epochs`. Pair with `--nnu-skip-prep` so the + preprocessed tree isn't `--clean`ed; the isolated Dataset221 is rebuilt + otherwise and a finished run's `checkpoint_latest.pth` is gone, so a + completed run + a bigger budget retrains B from scratch (interrupted runs + keep `checkpoint_latest.pth` and resume correctly). +- Verified 2026-10-04: A/C resumed epoch 1 → 3 (model restored, fresh + opt/sched for the old-format ckpt), re-run at 3/3 trained 0 epochs, B + continued from epoch 1 with LR restored mid-decay (0.00536 @ start of + epoch 2). -- **"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). +## Readiness notes for a full run -## 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 +- Budget knobs: `--train-epochs` (default 100), `--max-rows` (default 80% of + split_train ≈ 2856 for A/C and for B's `--nnu-cases`), `--max-eval-rows` + (default 10% of split_test ≈ 42), `--no-tta`. +- `--skip-train` is only usable for **A** (`runs/round0/best.pt`, trained + 2026-09-25 on the *previous* split — a few patients moved between + train/test in the rebuild, so its test dice carries a small leakage + caveat). No study checkpoints exist for B (Dataset210) or C + (`runs_monai/round0/best.pt`); they must be trained. +- GPUs: 3× 80 GB available; the benchmark uses one (`--gpu`). +- B's one-time cost scales with `--nnu-cases` (plan_and_preprocess ≈ 3.5 s/case + at 24 cases with 8 preprocess procs, plus training time at the planned + batch/patch). \ No newline at end of file diff --git a/scripts/02_build_m6_dataset.py b/scripts/02_build_m6_dataset.py index 9ab3675..87cff5e 100644 --- a/scripts/02_build_m6_dataset.py +++ b/scripts/02_build_m6_dataset.py @@ -1,12 +1,20 @@ """Build M6-2025 manifests: T1c volumes, labeled when GTV exists. -GTV lives in CT space; T1c is registered to CT (rigid, Mattes MI on the full -volumes -- no mask, no crop / pad; multi-start initial tilts because -inter-session head tilt reaches ~30 deg and tilted MI local minima can -outscore the true alignment). -The inverse transform then warps CT and GTV (nearest) into T1c native space. +GTV lives in CT space; the T1c is the fixed image and the CT is registered +to it with the best of three ANTsPy pipelines (src/registration_pipelines.py, +full volumes, no mask / crop / pad, each toolkit's native COM init): + ants_rigid : Rigid[0.25] (Mattes, 500x500x200, 4x2x1) + ants_affine : Affine[0.25] (Mattes, 2100x1200x1200x10, 6x4x2x1) + ants_syn : SyN (Affine[0.25] + SyN[0.2,3,0]) +The winner is the candidate with the highest Mattes-style MI between the +warped CT and the T1c (full volumes). All three return the +transform in the unified contract direction (T1c -> CT), so resampling the +CT / GTV onto the T1c grid with it places them in T1c native space -- the +TV is warped with SMD + linear (signed distance to the TV boundary, linearly +resampled, thresholded at the mean CT spacing halved; the validated recipe +from scripts/qa_label_interpolation.py). Acceptance requires the warped CT skull ring to be concentric with the MR head -silhouette (crop-independent check). +silhouette. Outputs: data/manifests/m6_labeled.jsonl {key, subject, date, img, label(warped GTV), source:'m6'} data/manifests/m6_unlabeled.jsonl {key, subject, date, img, label:None, source:'m6'} @@ -133,8 +141,8 @@ def _skull_ring_gap(warped_ct, ha): for a in np.linspace(0, 2 * np.pi, 72, endpoint=False): dy, dx = np.sin(a), np.cos(a) rr = np.arange(0.0, 160.0, 0.5) - yy = np.clip((cy0 + rr * dy).astype(int), 0, ha.shape[0] - 1) - xx = np.clip((cx0 + rr * dx).astype(int), 0, ha.shape[1] - 1) + yy = np.clip((cy0 + rr * dy).astype(int), 0, sl.shape[0] - 1) + xx = np.clip((cx0 + rr * dx).astype(int), 0, sl.shape[1] - 1) rh = rr[np.where(sl[yy, xx])[0]] rs = rr[np.where(sk[yy, xx])[0]] if len(rh) and len(rs): @@ -321,6 +329,112 @@ def register_ct_to_t1c(ct, t1c, threads=12): return tr, info, tissue_w +_QA = None + + +def _load_qa(): + """Import scripts/qa_label_interpolation.py (the validated SMD + linear + label-warp recipe: smd_threshold / make_tv_bin / warp_tv). Cached per + process (each worker process keeps its own copy).""" + global _QA + if _QA is None: + import importlib.util + p = os.path.join(os.path.dirname(os.path.abspath(__file__)), + "qa_label_interpolation.py") + spec = importlib.util.spec_from_file_location("qa_label_interp", p) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + _QA = m + return _QA + + +def register_ants_best(ct, t1c): + """Best-of-ANTs CT -> T1c registration with the T1c as the fixed image. + + Runs the three ANTsPy pipelines from src/registration_pipelines.py + (fixed = T1c, moving = CT, full volumes, no mask / crop / pad, each + toolkit's native COM init): + ants_rigid : Rigid[0.25] (Mattes, 500x500x200, 4x2x1) + ants_affine : Affine[0.25] (Mattes, 2100x1200x1200x10, 6x4x2x1) + ants_syn : SyN (Affine[0.25] + SyN[0.2,3,0]) + Each returns a transform in the unified contract direction (T1c -> CT), + so sitk.Resample(ct, t1c, T, ...) / sitk.Resample(tv, t1c, T, ...) place + the CT / TV concentric with the MR head -- no inversion (the T1c was the + fixed image). The winner is the candidate with the highest Mattes-style + MI between the warped CT and the T1c (_mi_value, full volumes). The + per-candidate skull-ring gap is still recorded and gates acceptance. + + Returns (transform, info, tissue_w) like register_ct_to_t1c. + """ + import time + from scipy import ndimage + from src import registration_pipelines as RP + from src.common import head_mask_from_image + t_start = time.time() + ct = sitk.Cast(ct, sitk.sitkFloat32) + t1c = sitk.Cast(t1c, sitk.sitkFloat32) + calls = { + "ants_rigid": lambda: RP.register_antspyx(ct, t1c), + "ants_affine": lambda: RP.register_ants_affine_variant(ct, t1c, "antspyx"), + "ants_syn": lambda: RP.register_ants_variant_deformable(ct, t1c, "ants_syn"), + } + ha = sitk.GetArrayFromImage(t1c) + cands = {} + for name, fn in calls.items(): + try: + T, meta = fn() + except Exception as e: # noqa + cands[name] = {"error": repr(e)} + print(f" {name}: FAILED {e!r}", flush=True) + continue + warped_ct = RP._warp_ct(ct, t1c, T) + mi = _mi_value(ha, warped_ct) + gap = RP._skull_gap(warped_ct, ha) + cands[name] = {"T": T, "meta": meta, "mi": mi, "gap": gap} + print(f" {name}: mi={mi:.4f} gap={gap} ({meta.get('seconds_reg')}s)", + flush=True) + ok = {n: c for n, c in cands.items() if "T" in c} + if not ok: + raise RuntimeError("all ANTs candidates failed: " + + repr({n: c.get("error") for n, c in cands.items()})) + best_name = max(ok, key=lambda n: ok[n]["mi"]) + T, meta, gap = ok[best_name]["T"], ok[best_name]["meta"], ok[best_name]["gap"] + + def _tp(x, y, z): + return np.array(T.TransformPoint((float(x), float(y), float(z)))) + + t0 = _tp(0, 0, 0) + mat = np.column_stack([_tp(1, 0, 0) - t0, _tp(0, 1, 0) - t0, _tp(0, 0, 1) - t0]) + scale = [round(float(np.linalg.norm(mat[:, i])), 4) for i in range(3)] + hm = head_mask_from_image(ha) + warped_ct = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, T, sitk.sitkLinear)) + posw = warped_ct[warped_ct > 0] + tissue_w = ((warped_ct > max(2.0, float(np.percentile(posw, 2)))) + if posw.size else np.zeros_like(warped_ct, bool)) + tissue_w = ndimage.binary_opening(tissue_w, iterations=2) + coverage = float(tissue_w[hm].sum() / max(hm.sum(), 1)) + info = { + "pipeline": "ants_best", + "fixed": "T1c", "moving": "CT", + "best": best_name, + "candidates": { + n: ({**c["meta"], "mi": round(c["mi"], 4), + "skull_ring_gap_mm": c["gap"], "selected": n == best_name} + if "T" in c else c) + for n, c in cands.items()}, + "affine_stage": {"matrix": [[round(float(v), 4) for v in row] for row in mat.tolist()], + "offset": [round(float(x), 3) for x in t0], "scale": scale, + "transform_type": type(T).__name__}, + "coverage_head_box": round(coverage, 3), + "skull_ring_gap_mm": gap, + "seconds_total": round(time.time() - t_start, 1), + } + print("reg done %.0fs best=%s mi=%.4f gap=%s cov %.2f" % ( + info["seconds_total"], best_name, ok[best_name]["mi"], gap, coverage), + flush=True) + return T, info, tissue_w + + def reg_worker(row): reg_dir = d("data/m6reg") key = row["key"] @@ -344,14 +458,17 @@ def reg_worker(row): f, indent=1) return key, False, 0 gtv = read_nii_img(row["label"]) - tr, info, tissue_w = register_ct_to_t1c(ct, t1c) + tr, info, tissue_w = register_ants_best(ct, t1c) # QA files in m6reg, everything in (oriented) MR coordinates: # _ct: CT warped to the MR grid (linear), for overlaying _tv # _mr: the oriented MR volume itself - # _tv: original TV warped to the MR grid (nearest) — the study label + # _tv: original TV warped to the MR grid (SMD + linear) — the study label # _reg.json: registration metrics + final image information - tv_mr = sitk.Resample(gtv, t1c, tr, sitk.sitkNearestNeighbor, 0.0) - z = sitk.GetArrayFromImage(tv_mr) + qa = _load_qa() + thr = qa.smd_threshold([float(s) for s in gtv.GetSpacing()])[0] + z = qa.warp_tv(qa.make_tv_bin(gtv), t1c, tr, "smd_linear", thr) + tv_mr = sitk.GetImageFromArray(z) + tv_mr.CopyInformation(t1c) nvox = int((z > 0.5).sum()) from scipy import ndimage tv_in_ct_frac = 1.0 @@ -369,6 +486,7 @@ def reg_worker(row): info.update({ "key": key, "subject": row["subject"], "date": row["date"], "tv_nvox": nvox, "tv_in_ct_frac": round(tv_in_ct_frac, 3), + "tv_smd_thr_mm": round(thr, 3), "alignment": align, "tv_accepted": accepted, "accept_reason": ("ok" if accepted else diff --git a/scripts/benchmark_pipelines.py b/scripts/benchmark_pipelines.py index 3834b33..6ed18fe 100644 --- a/scripts/benchmark_pipelines.py +++ b/scripts/benchmark_pipelines.py @@ -28,8 +28,11 @@ Protocol 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 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.json (+ _per_row.jsonl, _bench.png) @@ -81,6 +84,12 @@ def parse_args(): 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 benchmark checkpoint. A/C: load " + "runs_bench/{A,C}/best.pt (model + optimizer + scheduler when present) " + "and train the remaining epochs. 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 " @@ -208,7 +217,7 @@ def bench_A(args, train_rows, test_rows): 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) + 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) @@ -218,10 +227,26 @@ def bench_A(args, train_rows, test_rows): torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1), eta_min=3e-4 * 0.05)], milestones=[warmup]) + ckpt = os.path.join(ckpt_dir, "best.pt") + start_epoch = 0 + 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: + 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(args.train_epochs): + for ep in range(start_epoch, args.train_epochs): if hasattr(dl.sampler, "set_epoch"): dl.sampler.set_epoch(ep) loss_sum = n = 0.0 @@ -245,11 +270,11 @@ def bench_A(args, train_rows, test_rows): 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_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 - ckpt = os.path.join(ckpt_dir, "best.pt") - torch.save({"model": model.state_dict(), "epoch": args.train_epochs}, ckpt) + torch.save({"model": model.state_dict(), "epoch": args.train_epochs, + "optimizer": opt.state_dict(), "scheduler": sched.state_dict()}, 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())) @@ -326,10 +351,26 @@ def bench_C(args, train_rows, test_rows): eta_min=3e-4 * 0.05)], milestones=[warmup]) loss_f = WeightedDiceCELoss() + ckpt = os.path.join(ckpt_dir, "best.pt") + start_epoch = 0 + 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: + 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(args.train_epochs): + for ep in range(start_epoch, args.train_epochs): if hasattr(dl.sampler, "set_epoch"): dl.sampler.set_epoch(ep) loss_sum = n = 0.0 @@ -350,11 +391,11 @@ def bench_C(args, train_rows, test_rows): 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_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 - ckpt = os.path.join(ckpt_dir, "best.pt") - torch.save({"model": model.state_dict(), "epoch": args.train_epochs}, ckpt) + torch.save({"model": model.state_dict(), "epoch": args.train_epochs, + "optimizer": opt.state_dict(), "scheduler": sched.state_dict()}, 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())) @@ -432,10 +473,21 @@ def bench_B(args, train_rows, test_rows): 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 + 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(N.train_cmd(1), os.path.join(d("logs/nnu"), "bench_train.log"), env=env, retries=0) + N.run(bcmd, 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