From e7721760f23314fe17f8026c2b27229a9f185eb9 Mon Sep 17 00:00:00 2001 From: Furen Xiao Date: Thu, 8 Oct 2026 05:28:15 +0800 Subject: [PATCH] Enhance MONAI training scripts and models for Pipeline B and C - 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. --- .gitignore | 2 + BENCHMARKING.md | 500 +++++++++-------- README.md | 51 +- REGISTRATION_BENCHMARK_AFFINE.md | 4 +- REGISTRATION_BENCHMARK_DEFORMABLE.md | 6 +- REGISTRATION_BENCHMARK_RIGID.md | 16 +- scripts/bench_registration.py | 18 +- scripts/benchmark_pipelines.py | 694 ++++++++++++++++++++---- scripts/run_b_100ep_stages.sh | 37 ++ scripts/run_c_bloss_100ep_after_b.sh | 19 + scripts_monai/02_monai_train.py | 133 +++-- scripts_monai/03_monai_pseudo_label.py | 11 +- scripts_monai/04_monai_eval_test.py | 9 +- scripts_monai/05_monai_run_iterative.py | 25 +- scripts_monai/monai_common.py | 178 +++++- src/augment_b.py | 149 +++++ src/dataset.py | 32 +- src/losses.py | 69 ++- src/training.py | 19 +- src/unet3d.py | 58 +- 20 files changed, 1584 insertions(+), 446 deletions(-) create mode 100644 scripts/run_b_100ep_stages.sh create mode 100644 scripts/run_c_bloss_100ep_after_b.sh create mode 100644 src/augment_b.py diff --git a/.gitignore b/.gitignore index ba2b644..b41c87e 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,5 @@ __pycache__/ elastix/ niftyreg/NiftyReg-Ubuntu-v2.1.1.zip niftyreg/NiftyReg-Ubuntu-v2.1.1/ + +.kilo/ \ No newline at end of file diff --git a/BENCHMARKING.md b/BENCHMARKING.md index eb30b4f..23843c6 100644 --- a/BENCHMARKING.md +++ b/BENCHMARKING.md @@ -6,258 +6,330 @@ Benchmark of the three segmentation pipelines — A in-house Unet3D, B nnU-Net v the script docstring; this file records what has been run and what the machine currently supports. -## Data chain status (2026-10-04) +## Data (2026-10-04 curation) -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: +M6 labeled set rebuilt 2026-10-04 (`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 → 1312 labeled volumes, all 1541 +processed (`data/proc/{,_label}.nii.gz`, 1 mm). Row-level curation lives +in `src/exclusions.py` (applied at `scripts/make_tv_slices.py`, re-applied +after split by `05_build_splits`): total TV > 1e6 mm3 (1 row), +single-component TV < 100 mm3 (109), non-brain images (33, all ntuh +spine/neck/torso; 6 skull-base/posterior-fossa cases kept after expert +review — protocol: classify at the mid-cranial slice, not just the TV-max +slice). Audit trail: `data/manifests/excluded.jsonl` (143 rows); +pre-exclusion manifests in `data/manifests/backup_20261004/`. -- `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% +Current splits (seed 0): **4057** labeled rows → train 3447 / val 205 / test +405 (3771 subjects); unlabeled pool 13280. Tumor volume: median 3212 mm3, +p95 23.4k, max 163k (m6 meningioma), zero_frac 0.12%. Benchmark budget: +**2750** train rows (80% of split_train), **40** eval rows (10% of +split_test), TTA on. -## Dataset exclusions (2026-10-04) +## Environment -Row-level curation lives in `src/exclusions.py` and is applied at the **first -stage** by `scripts/make_tv_slices.py` (right after the rows are fetched from -the sources, before any split): it classifies every labeled volume and writes -its TV-slice QA figure to `data/tv_slices/inclusion/` (kept) or -`data/tv_slices/exclusion/` (excluded, reason in the title). `05_build_splits` -re-applies the same rules **after** split assignment (kept rows retain their -split; the subject-level shuffle boundaries are unchanged): +- `monai==1.6.1` in `longitudinal-313` (pip `--no-deps`, torch untouched). +- **B's nnU-Net tree is on local disk, not CIFS**: `nnu_common.nnu_root()` + honors `NNU_ROOT=/home/xfr/longitudinal_nnu` (local ZFS) and B's pred I/O + honors `BENCH_B_DIR=/home/xfr/longitudinal_benchB` — the 2026-10-05 CIFS + flapping (`OSError [Errno 112]`, stale `rmtree` → `FileExistsError`) had + killed four B attempts. CIFS is then touched only for the one-time raw + reads from `/mnt/pve` and small log/JSON appends. Hardening: + `nnu_common.rmtree_verified()`, plan retries 2→5, full `traceback` in + failure JSONs. -- total TV > 1e6 mm3: **1** row (ntuh_6019528, label outlines the full skull) -- single target (1 connected component) with TV < 100 mm3: **109** rows -- non-brain images: **33** rows (all ntuh — spine/neck/torso CyberKnife - cases whose series names are indistinguishable from brain T1c; identified - by visual review of the TV-slice figures, now in `data/tv_slices/exclusion/`) +## Headline correction: A's low dice was an inference TTA bug (2026-10-06) -Audit trail: `data/manifests/excluded.jsonl` (143 rows, reason + split); -pre-exclusion manifests in `data/manifests/backup_20261004/`. New dataset: -**4057** labeled rows → train 3447 / val 205 / test 405 (3771 subjects); -unlabeled pool unchanged (13280); volume stats now median 3212 mm3, p95 -23.4k, max 163k (m6 meningioma), zero_frac still 0.12% (5 rows, e.g. -ntuh_3436472 ×3). +A's persistent low test dice (pmax median ~0.25, "under-confidence state") in +the 2026-10-05/06 10–100-epoch runs was a bug in +`src/training.sliding_window_probs`, not a training problem: the old code +flipped the **whole volume** and windowed at the **original grid positions**, +so each region averaged its own model response with its mirror regions' +responses — for an off-mirror target the four-view average collapses to ~1/4 +of the identity-view probability. C's flip→window→flip-back scheme was +already correct, hence unaffected. Evidence (A 100ep model): p = 0.998 at +the tumor voxel with TTA off (top-10 activations 10/10 inside the tumor, +mean p inside 0.91 vs 0.004 outside), pmax = 0.2498 with TTA on; a +BN→InstanceNorm conversion changed nothing (0.2498 → 0.2499). Fix: flip the +**patch** (`patch = tf(vol[sl])`). -After the initial 39-row non-brain list, expert inspection overturned 6 -(skull-base / posterior-fossa brain targets that look like cervical spine at -the TV-max slice: ntuh_5902230, ntuh_5231609, ntuh_5083686, ntuh_5157194, -ntuh_5303758, ntuh_6373987). The remaining 33 were re-audited at the -mid-cranial slice (widest head-mask axial) — all still non-brain. Review -protocol now recorded in `src/exclusions.py`: classify at full resolution -on the mid-cranial slice, not just the TV-max slice. +All A checkpoints re-scored on 40 held-out rows +(`results/benchmark_full/A_reeval_fixed_tta.json`): -## Smoke run (2026-10-04) +| A checkpoint | dice (old TTA) | dice (fixed) | pmax median | rows > 0.001 | +|---|---|---|---|---| +| 100ep nn-loss (AdamW, BN) | 0.0047 | **0.5313** | 0.9994 | 30/40 | +| 10ep baseline, uniform crops, std loss | 0.0026 | **0.5102** | 1.0000 | 31/40 | +| 20ep nn-loss (AdamW, BN) | 0.0000 | 0.3644 | 0.9532 | 26/40 | +| 20ep std-loss + CE w=500 | 0.0076 | 0.1466 | 0.9982 | 33/40 | -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): +The 2026-10-05 fg-crop / CE-weight / loss-parity investigations were chasing +a symptom of this eval bug — the plain uniform-crop 10ep baseline (0.5102) +already beats everything else on record at that budget. (Caveat: that +checkpoint predates the current split by a few moved patients — same leakage +caveat as `--skip-train`.) -``` -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 -``` +## What pipeline B actually does (2026-10-06 audit) -| 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 | +From the installed nnunetv2 2.8.1 source + B's `fold_0/debug.json`: -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`. +- **Loss**: `DC_and_CE_loss({'batch_dice': True, 'smooth': 1e-5, 'do_bg': + False})` = CE + tumor-class-only batch-level soft dice, `loss = CE − dice` + (`src.losses.b_dc_and_ce` matches nnunetv2's classes to max abs diff 0.0). +- **Deep supervision**: 4 seg outputs at 1x/1/2/1/4, weights 4/7, 2/7, 1/7 + (raw [1, 1/2, 1/4, 1/8] with the bottleneck scale zeroed, renormalized), + nearest-neighbor downsampled targets. +- **Optimizer**: SGD lr 1e-2, momentum 0.99, nesterov, weight_decay 3e-5; + PolyLR exponent 0.9 on the global epoch index; grad clip 12. +- **Network**: PlainConvUNet, InstanceNorm3d(affine, eps 1e-5), + LeakyReLU(0.01), conv bias, init kaiming_normal_(a=0.01) + zero bias. +- **Data**: resampled to spacing [1.0, 0.586, 0.586], ZScore (no mask), + patch 64×192×192, batch 2, 250 iterations/epoch, 33% fg patch oversampling. +- **Augmentation** (per patch): rotation ±30° p=0.2/axis; zoom 0.7–1.4 + p=0.2; Gaussian noise p=0.1; blur p=0.2; brightness/contrast p=0.15; + low-resolution sim p=0.25; gamma p=0.1 + p=0.3; mirror p=0.5; no elastic. +- **"EMA"**: an EMA (0.9/0.1) of the per-epoch val pseudo-dice, used only to + select `checkpoint_best.pth` — not a weight EMA. +- **Epoch semantics**: B's epoch = 250 iters × batch 2 = 500 patches; A/C's + epoch = full pass = 2750 patches (B @100ep ≈ A/C @~18ep in data seen). -### Environment fixes made during the check +## B extended to 100 epochs (2026-10-06, complete 14:59 UTC) -- `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. +Staged extension from the 20ep checkpoint via +`scripts/run_b_100ep_stages.sh`: `--resume --nnu-skip-prep +--train-epochs 50` then `100`, scoring `checkpoint_final.pth` (the +last-trained model, A/C-comparable) via `--nnu-ckpt`; stage 3 additionally +scored `checkpoint_best.pth` with `--skip-train`. -## Full-dataset 10-epoch run (2026-10-05, updated dataset) +| dice @20 | dice @50 | dice @100 (final) | dice @100 (best) | +|---|---|---|---| +| 0.2927 | 0.3333 | **0.3912** | 0.3912 (== final) | -`--train-epochs 10`, default budgets on the current curated splits: **2750** -train rows = 80% of split_train (3438 labeled), **40** eval rows = 10% of -split_test (404), TTA on. The 2026-10-04 interim run (pre-exclusion budget -2856/42: A 0.0008, C 0.2119) is superseded; its JSONs are kept in -`results/backup_20261004/` (smoke B checkpoint in -`nnu/results/backup_20261004_Dataset221_smoke/`). +B improved monotonically; 32/39 rows > 0.001 @100 (vs 29/39 @20); +138.5 s/epoch stage pace. Artifacts: `results/benchmark_full/ +benchmark_full_B_{50ep,100ep,100ep_best}.json`. -``` -python scripts/benchmark_pipelines.py --pipelines C --gpu 2 --train-epochs 10 \ - --out results/benchmark_full_C.json # done -python scripts/benchmark_pipelines.py --pipelines A --gpu 0 --train-epochs 10 \ - --a-fg-ce-weight 500 --out results/benchmark_full_A.json # done -# A baseline (uniform crops) on this dataset: dice 0.0026, see table; -# checkpoint runs_bench/A/best_baseline_10ep_uniform_crop.pt -NNU_ROOT=/home/xfr/longitudinal_nnu BENCH_B_DIR=/home/xfr/longitudinal_benchB \ -python scripts/benchmark_pipelines.py --pipelines B --gpu 1 --train-epochs 10 \ - --out results/benchmark_full_B.json # done (local-disk NNU_ROOT, see below) -``` +## B-aligned A/C (C's default; A's opt-in `--b-aligned`) -| pipeline | train s/epoch | wall | tr VRAM | s/vol (TTA) | in VRAM | params | weights | dice @10ep | -|---|---|---|---|---|---|---|---|---| -| C MONAI UNet(16→128) | 256.7 (916 steps) | 43 min | 0.4 GB | 0.37 ± 0.06 | 0.5 GB | 1.19 M | 14.3 MB | **0.2162** | -| A Unet3D(16, depth 4), improved (fg crops + CE w=500) | 232.3 (916 steps) | 41 min | 2.0 GB | 3.60 ± 0.76 | 0.5 GB | 1.46 M | 17.6 MB | **0.0369** | -| A baseline (uniform crops, unweighted loss) | 216.6 (916 steps) | 45 min | 2.0 GB | 3.33 ± 0.86 | 0.5 GB | 1.46 M | 17.6 MB | **0.0026** | -| B nnU-Net 3d_fullres | 331.5 (incl. 275-case val/epoch) | ~2.5 h (raw 6 m, prep 62 m, train 55 m, infer 14 m) | 9.3 GB | 21.84 | 1.8 GB | 30.8 M | 246 MB | **0.1424** (hard 0.1424) | +Training C with `benchmark_pipelines.py` (no config flags) — and training A +with `--b-aligned` — uses B's full protocol on top of the backbone (A: +`Unet3D(16,4)` + LeakyReLU; C: MONAI `UNet(16→128)`), keeping the benchmark +data budget (2750 train / 40 eval rows, 96³, batch 3, bf16): +InstanceNorm3d(affine) + kaiming-normal(0.01) init, B's true loss +(`b_dc_and_ce`) with deep supervision (4/7, 2/7, 1/7 on 1x/1/2/1/4, NN +downsampled targets), SGD 1e-2 (0.99, nesterov, wd 3e-5) + PolyLR(0.9), B's +augmentation + Z-score + 33% fg sampling (`src/augment_b.py`, patch-level, +shared by A and C), grad clip 12, and B-style per-epoch val pseudo-dice (50 +val patches, EMA 0.9/0.1) selecting `runs_bench/{A,C}/best_ema.pt`. Code: +`src/unet3d.py` (norm/deep-sup/init), `src/losses.py` (b_dc_and_ce, +deep_sup_loss), `src/dataset.py` (zscore/b_aug), +`scripts_monai/monai_common.py`, plus the TTA fix above. A's default is its +own A-native config (see the ablation section below); opt-outs: `--legacy` +(pre-2026-10-07 config, approx nn-loss default; `best_legacy.pt`), +`--b-loss-only` (B's loss, rest legacy; `best_bloss.pt`). -C: 25/39 timed test vols with dice > 0.001 (max 0.7392); loss 0.403 → 0.182. -10 epochs is far below convergence for these sparse targets (median tumor -3.2 k mm³); A/C numbers at this budget are intermediate. Checkpoints at -`runs_bench/{A,C}/best.pt` (extend with `--resume --train-epochs `). +From-scratch 100-epoch runs (40 held-out test rows, TTA on; A/C epoch = full +pass, B epoch = 500 patches): -### A: why dice ≈ 0, and the fix (2026-10-05) +| run (100ep, from scratch) | dice @5 | @10 | @20 | @50 | @100 | +|---|---|---|---|---|---| +| **A B-aligned** | 0.2017 | 0.3176 | 0.4368 | 0.4998 | **0.5417** | +| **A all-default** (`--legacy --std-loss`) | 0.3161 | 0.4578 | 0.4249 | 0.5089 | **0.5508** (A record; see ablation below) | +| **C B-aligned** | 0.0587 | 0.1748 | 0.4073 | 0.4396 | **0.4002** | +| **C all-default** (`--legacy --std-loss`) | 0.2025 | 0.2209 | 0.0974 | 0.2872 | **0.2083** (see C ablation below) | +| C b-loss-only (true B loss, rest legacy) | 0.2051 | 0.3267 | 0.2824 | 0.1710 | 0.2523 | +| A legacy (BN, AdamW; re-scored, fixed TTA) | — | — | 0.3644 | — | 0.5313 | +| C legacy (BN, AdamW, approx nn-loss) | 0.1889 | 0.2562 | 0.3358 | 0.1798 | 0.2649 | +| B nnU-Net (reference) | — | — | 0.2927 | 0.3333 | 0.3912 | -The A baseline predicts all-background (38/39 test rows dice exactly 0.0). -Root-cause work: +- **Both B-aligned models beat B (and their legacy counterparts) at every + point ≥20ep.** A: 0.5417 vs 0.5313 legacy — the TTA fix was the big A win, + the B protocol adds a smaller one on top. C: 0.4002 vs 0.2649 legacy + (+0.14) and the legacy mid-training collapse is GONE — monotone + 0.4073 → 0.4396 → 0.4002 (legacy: 0.3358 → 0.1798 → 0.2649). C's slow + start (0.0587 @5) is the SGD/PolyLR + instance-norm warm-up. +- **A 100→200ep extension** (`--resume --train-epochs 200 --eval-at "150 +200"`, complete 2026-10-07 05:14 UTC, 139.9 s/epoch): @150 0.5072, @200 +0.5397 — no improvement over the @100 B-aligned (0.5417); A is converged. + The resume restarts PolyLR at a higher LR over the 200-ep total, costing + a @150 dip; val pseudo-dice EMA kept rising (0.3622 @200) while test dice + plateaued — the val-EMA proxy over-estimates later epochs. The @200 final + save overwrote `best_bal.pt`/`best.pt` with the @200 model; it is 0.002 + from the @100 B-aligned (within noise), so the @200 weights on disk are + the best A B-aligned model. +- **The loss alone does NOT fix C's instability**: b-loss-only (B's exact + loss, legacy BN/AdamW/no deep-sup/no B aug) repeats the legacy dip — + statistically the same trajectory as C-legacy. The full package (instance + norm + deep supervision + SGD/PolyLR + B augmentation + Z-score) is what + stabilizes C. +- **B's best-checkpoint heuristic does not transfer**: for the B-aligned + models the EMA-selected ckpt is WORSE than the final model — A's EMA-best + happens to be epoch 100 (= final), but C's EMA-best (epoch 97, val-ema + 0.2442) scores 0.2829 on test vs 0.4002 for the epoch-100 model. Use the + final checkpoint for A/C. +- pmax saturation is immediate under B-aligned training (A: 0.999 @50; no + under-confidence at any point — the old pmax ~0.25 was purely the TTA bug). +- Pace: A-bal ~150 s/epoch, C-bal ~175 s/epoch (patch-level B augmentation; + the initial volume-level C attempt ran ~600 s/epoch and was restarted — B's + SpatialTransform acts on the sampled patch, so patch-level is the faithful + and fast variant). -1. **Uniform random 96³ crops starve the loss of foreground.** A 96³ patch is - 885 k voxels; the median tumor is ~4.6 k mm³ (0.02–0.05% of a head volume). - Simulated uniform crops contain the whole blob in only 8–37% of patches and - a *median* crop often has 0 tumor voxels. Pipeline C's - `RandCropByPosNegLabeld(pos=1, neg=1)` centers half its patches on the - tumor — ~4× more foreground supervision per epoch. -2. **Fg-aware crops alone are not enough: A under-confides.** With the crop - protocol fixed (`src/dataset.py`), A's loss matches C's at every epoch - (0.193 vs 0.182 @10ep) yet p_tumor maxes at 0.43 (C: ~1.0): it localizes - (top-k dice 0.48 on best rows) but never crosses the 0.5 threshold, so - dice@0.5 stays 0. -3. **Fix: class-weighted CE** — tumor-class weight `w` in the CE part of the - per-sample loss (`--a-fg-ce-weight`, `src/losses.py`). 3-epoch probes on - this dataset: w=20 → pmax median 0.33, dice 0.0; w=100 → 0.41, dice - 0.0002; w=500 → 0.62, dice 0.0188 (34/39 rows > 0.5). Official 10-epoch - w=500 run: **dice 0.0369** (vs 0.0026 baseline, 14×; 8/39 rows > 0.001, - best 0.555, pmax median 0.55, 29/39 rows > 0.5) — A now segments, but - still trails C (0.2162) at this budget; its loss is still descending at - epoch 10 (0.78 → 0.51), i.e. under-converged vs C (0.41 → 0.18), so the - gap should narrow with more epochs. +Artifacts: `results/benchmark_full/benchmark_full_{A,C}_bal_100ep.json` +(+ `_per_row.jsonl`, progress PNGs), `benchmark_full_A_bal_200ep.json`; +checkpoints `runs_bench/C/best.pt` (C @100 B-aligned 0.4002, C's record) / +`best_ema.pt` (EMA-selected); A's B-aligned models: `runs_bench/A/ +best_bal.pt` (@200, 0.5397 — the best A B-aligned model) / +`best_ema_bal.pt`. -Changes: `src/dataset.py` (pos/neg crop sampling `pos_neg=0.5`, per-worker -RNG reseed), `src/losses.py` (`soft_dice_fg`, `fg_ce_weight`), benchmark flags -`--a-fg-dice` / `--a-fg-ce-weight`, per-row `pmax` in A's `dice_per_row`, -full traceback in failure JSONs. A/B probes: `results/bench_A_*probe*.json`. +## A ablation: all-default config — now A's default (2026-10-07) -### B: CIFS failures and local-disk NNU_ROOT (2026-10-05) +With the TTA fixed, A B-aligned (0.5417 @100ep) vs A legacy (0.5313 @100ep, +approx nn-loss) differs by many components at once. This run trained A with +**every hyperparameter at A's default, none nnU-derived** — the fully +A-native baseline, whose mid-training trajectory was never recorded (the old +100ep legacy run's @5/@10/@50 evals were scored with the buggy TTA): -The CIFS `/mnt/b4` mount dropped I/O repeatedly on 2026-10-05 (NAS flapping: -`OSError [Errno 112] Host is down`, `RuntimeError: Error while setting the -slice`, one silent `rmtree` failure that made the raw rebuild collide with -stale entries → `FileExistsError`). Four B attempts died (07:52 plan&prep; -11:12 raw EEXIST; ~12:40 / ~13:50 / ~14:00 plan&prep). Fixes: +BatchNorm + PyTorch default init (no kaiming), ReLU, per-sample +0.5CE+0.5mean-dice loss (no deep supervision — the `--legacy` default loss +is the nnU-style approx `nn_dc_and_ce`, so the run used `--std-loss` to make +the loss A-native too), AdamW 3e-4 (wd 1e-4) + warmup/cosine, grad clip 5.0, +percentile [0,1] input (no Z-score), legacy augmentation +(flip/rot90/brightness/noise), 50% fg-centered crops (not B's 33%). +Benchmark constants unchanged for comparability: 2750 train / 40 eval rows, +batch 3, 96³, bf16, TTA on, eval at 5/10/20/50. (Run as +`--legacy --std-loss`; since the promotion below it is A's plain default.) -- `nnu_common.rmtree_verified()` + defensive per-link unlink in - `make_raw_dataset` (EEXIST on rebuild) -- failure JSONs now carry a full `traceback` key -- plan step retries raised 2 → 5 (`N.run(..., retries=5)`) -- **B's nnU-Net tree moved off CIFS**: `nnu_common.nnu_root()` honors - `NNU_ROOT` (set to `/home/xfr/longitudinal_nnu`, local ZFS ~1.8 T free) and - the benchmark's B pred I/O honors `BENCH_B_DIR` - (`/home/xfr/longitudinal_benchB`). CIFS is then touched only for the - one-time raw image reads from `/mnt/pve` and small log/JSON appends. - The pre-`NNU_ROOT` attempt's half-built preprocessed tree under - `nnu/preprocessed/` is stale — the local tree is the one being used. +**Complete 2026-10-07 ~09:45 UTC** (persistent background process on GPU 0, +shared with co-tenant jobs; 125.1 s/epoch, train VRAM 1.98 GB, 2.78 s/vol +TTA; log `logs/bench_A_legacy_std_100ep.log`): -The local-disk run succeeded on attempt 1 (2026-10-05 14:07→16:38 UTC): -raw build 6.3 min, plan&preprocess 62 min (140 GB preprocessed), 10 epochs -in 55 min (epoch times 101–167 s train + 275-case validation — -`nnUNet_compile is enabled`, vs the Oct-4 eager-mode pace of ~2.4 h/epoch), -inference 21.8 s/vol TTA. Val pseudo dice 0.091 → 0.459, best EMA 0.285; -test dice 0.1424 (29/39 rows > 0.001, max 0.675). Checkpoints in -`/home/xfr/longitudinal_nnu/results/Dataset221_T1C_BENCH/...fold_0/` -(copy to `nnu/results/` for the record if local space is reclaimed). +| dice (40 held-out rows, TTA on) | @5 | @10 | @20 | @50 | @100 | +|---|---|---|---|---|---| +| **A all-default** (`--legacy --std-loss`) | 0.3161 | 0.4578 | 0.4249 | 0.5089 | **0.5508** (A record) | +| A B-aligned | 0.2017 | 0.3176 | 0.4368 | 0.4998 | 0.5417 | -## Extended to 20 epochs (2026-10-05/06) +- **A's own defaults beat B's protocol at every eval point** — final + 0.5508 (30/40 rows > 0.001, pmax median 0.9989) vs 0.5417 B-aligned; + 0.5508 is the best result on record for any pipeline (B 0.3912, C + 0.4002). Loss 0.386 → 0.161, flat from ~epoch 40. The earlier "B + protocol adds +0.01 for A" (0.5417 vs 0.5313) was against the legacy + config with the nnU-style approx loss: the only difference between the + 0.5313 model and this run is the loss, so A's per-sample 0.5CE+0.5dice + alone is worth +0.0195 — and the full B package (IN + deep sup + + SGD/PolyLR + B aug + Z-score + 33% fg) nets −0.009 on this backbone. + Contrast C, where the same B package was worth +0.14 (0.4002 vs 0.2649) + and remains C's record and default. +- pmax saturation is immediate (0.81 @5 → 0.9998 @10) — no + under-confidence at any point; with the TTA fix, A's A-native config is + the best A result on record. +- Checkpoint: the 100ep all-default model (17.6 MB). The prior legacy 100ep + model (0.5313) had already been lost to a smoke-run overwrite; its number + is recorded in `results/benchmark_full/A_reeval_fixed_tta.json`. -All three continued from their 10-epoch checkpoints to a 20-epoch total -(`--train-epochs 20` = total budget; 10-epoch JSONs snapshotted as -`results/benchmark_full_{A,B,C}_10ep.json`): +**Promoted to A's default (2026-10-07).** Running `benchmark_pipelines.py` +with no config flags now trains A with this A-native config and writes +`runs_bench/A/best.pt` (the 0.5508 model is there); `--b-aligned` opts A +into B's protocol (`best_bal.pt`), `--legacy` gives the pre-2026-10-07 A +config (approx nn-loss; `best_legacy.pt`), `--b-loss-only` the B-loss +ablation (`best_bloss.pt`), and deviating A knobs (`--a-leaky`, +`--a-fg-*`) write `best_alt.pt` so the record checkpoint is never +clobbered. C's default is unchanged (B-aligned, `runs_bench/C/best.pt`). +Verified: all four A modes + C `--legacy` smoke-tested (1 epoch each, +config/loss/ckpt recorded in the JSONs), and a no-op +`--resume --train-epochs 100` on the new default re-scored 0.5508 exactly +(30/40 rows, pmax 0.9989). -``` -python scripts/benchmark_pipelines.py --pipelines A --gpu 0 --resume --train-epochs 20 \ - --a-fg-ce-weight 500 --out results/benchmark_full_A.json -python scripts/benchmark_pipelines.py --pipelines C --gpu 2 --resume --train-epochs 20 \ - --out results/benchmark_full_C.json -NNU_ROOT=/home/xfr/longitudinal_nnu BENCH_B_DIR=/home/xfr/longitudinal_benchB \ -python scripts/benchmark_pipelines.py --pipelines B --gpu 1 --resume --nnu-skip-prep \ - --train-epochs 20 --out results/benchmark_full_B.json -``` +## C ablation: all-default config (2026-10-07) -A/C resume from `runs_bench/{A,C}/best.pt` (model + optimizer + scheduler, -epochs 11–20). B: the completed run had deleted `checkpoint_latest.pth` -(see Resuming training), so it was restored from `checkpoint_final.pth` -(`current_epoch: 10`) before `--c` + `--nnu-skip-prep` — optimizer state and -LR schedule restored (initial validation 0.1253). `bench_B`'s resume -s/epoch now divides by epochs actually trained (reads `current_epoch` from -the checkpoint). +The mirror of A's ablation on the MONAI backbone: C trained with every +hyperparameter at C's own default (its original study protocol from +`scripts_monai/02_monai_train.py`, nothing nnU-derived) — BatchNorm +`UNet(16→128)` with MONAI default init, LeakyReLU (C's own default act), +per-sample 0.5CE+0.5mean-dice loss (WeightedDiceCELoss, no deep +supervision), AdamW 3e-4 (wd 1e-4) + warmup/cosine, grad clip 5.0, +percentile [0,1] input (no Z-score), legacy augmentation +(flip/rot90/brightness/noise), 50% fg-centered crops. Benchmark constants +unchanged: 2750 train / 40 eval rows, batch 3, 96³, bf16, TTA on, eval at +5/10/20/50. Run as `--legacy --std-loss` (C's default remains B-aligned). -| pipeline | dice @10ep | dice @20ep | Δ | median | max | rows > 0.001 | -|---|---|---|---|---|---|---| -| B nnU-Net 3d_fullres | 0.1424 | **0.2927** (hard 0.2927) | +0.150 | 0.289 | 0.839 | 29/39 | -| C MONAI UNet(16→128) | 0.2162 | **0.1621** | −0.054 | 0.126 | 0.680 | 31/39 | -| A Unet3D improved (fg crops + CE w=500) | 0.0369 | **0.0076** | −0.029 | 0.000 | 0.196 | 6/39 | +**Complete 2026-10-07 ~18:31 UTC** (persistent background process on GPU 0; +254.3 s/epoch, train VRAM 0.48 GB, 0.39 s/vol TTA; 1-epoch smoke first, +config recorded in `benchmark_full_C_legacy_std_smoke1.json`): -Resume window: A 229.2 s/epoch, C 253.9 s/epoch, B 280.1 s/epoch -(10 epochs each, same data budget). B was far from convergence at 10 epochs -and improved a lot; C's loss was flat (0.184) and its extra epochs slightly -degraded test dice; the A w=500 config oscillates (loss 0.49–0.55, pmax -median 0.55 → 0.49, rows >0.5: 29 → 19), so its 10-epoch point is its best -under this config. At 20 epochs **B leads (0.2927)**. Checkpoints in -`runs_bench/{A,C}/best.pt` and B's fold_0 now hold the 20-epoch models — the -10-epoch models were overwritten (JSONs preserved). B: 21.9 s/vol TTA, -9.8 GB train VRAM. +| dice (40 held-out rows, TTA on) | @5 | @10 | @20 | @50 | @100 | +|---|---|---|---|---|---| +| **C all-default** (`--legacy --std-loss`) | 0.2025 | 0.2209 | 0.0974 | 0.2872 | **0.2083** | +| C B-aligned (record) | 0.0587 | 0.1748 | 0.4073 | 0.4396 | **0.4002** | +| C legacy (approx nn-loss) | 0.1889 | 0.2562 | 0.3358 | 0.1798 | 0.2649 | +| C b-loss-only | 0.2051 | 0.3267 | 0.2824 | 0.1710 | 0.2523 | + +- **The non-B package is unstable for C whatever loss it carries**: this + run dips @10→@20 (0.2209 → 0.0974) and ends 0.2083 — the worst of the + three legacy-family C configs (nn-loss 0.2649, b-loss-only 0.2523) — even + though the TRAIN loss is smooth (0.35 → 0.15, no collapse; min 0.153 @88): + the test dice oscillates, the same family as C-legacy's + 0.3358 → 0.1798 @20→@50 dip. B-aligned beats it at every point ≥20ep + (+0.31 @20, +0.15 @50, +0.19 @100). +- **Mirror image of A**: on A's backbone the all-default config is A's best + (0.5508; every A variant sits in 0.53–0.55), on C's backbone it is C's + worst (0.2083; every non-B C variant sits in 0.21–0.27). The B package + (instance norm + deep supervision + SGD/PolyLR + B augmentation + Z-score) + is not a preference — it is what stabilizes the MONAI UNet. And A's + winning loss switch (per-sample 0.5CE+0.5dice, +0.0195 for A) does not + transfer: for C it is the worst legacy-family result. +- Checkpoint: the @100 model (14.4 MB) is `runs_bench/C/best_legacy.pt` — + it overwrote the prior C legacy nn-loss model (0.2649; number recorded in + `benchmark_full_C_nnloss_100ep.json`). C's record B-aligned model + (`best.pt`, 0.4002) is untouched. + +Artifacts: `results/benchmark_full/benchmark_full_C_legacy_std_100ep.json` +(+ `_per_row.jsonl`, progress PNG). ## Resuming training (`--resume`) -All three pipelines can continue from their latest benchmark checkpoint: +`--train-epochs` is the **total budget** (not epochs added). -``` -python scripts/benchmark_pipelines.py --pipelines "A B C" --resume \ - --train-epochs 200 --nnu-skip-prep -``` - -- `--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 +- **A / C**: resume from the latest checkpoint of the **same config mode** + (model + optimizer + scheduler when present; A's B-aligned PolyLR is + stateless — the LR is recomputed from the epoch index over the new total, + same staged convention as B's extensions). A: `best.pt` (A-native default) + / `best_bal.pt` (`--b-aligned`) / `best_legacy.pt` (`--legacy`) / + `best_bloss.pt` (`--b-loss-only`); C: `best.pt` (B-aligned default) / + `best_legacy.pt` / `best_bloss.pt`. Checkpoints written before + 2026-10-04 carry model + epoch only and fall back to fresh + optimizer/scheduler. A run is a no-op train (eval only) when the + checkpoint already reached the total. +- **B**: appends nnU-Net's `--c` (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). A finished run CAN - still be resumed by copying `checkpoint_final.pth` → `checkpoint_latest.pth` - first (same state dict; used for the 20-epoch extension above). -- 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). + `NNUNet_pl_EPOCHS = --train-epochs`; pair with `--nnu-skip-prep` so the + preprocessed tree isn't `--clean`ed. Completed runs delete + `checkpoint_latest.pth` — copy `checkpoint_final.pth` → + `checkpoint_latest.pth` first to resume (same state dict; used for B's + 20→50→100ep stages). -## Readiness notes for a full run +## Notes -- Budget knobs: `--train-epochs` (default 100), `--max-rows` (default 80% of - split_train = 2750 for A/C and for B's `--nnu-cases` on the current curated - splits; was 2856 pre-exclusion), `--max-eval-rows` (default 10% of - split_test = 40), `--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. +- Per-pipeline defaults (since 2026-10-07): A = A-native all-defaults (the + 0.5508 @100ep record; B's protocol is `--b-aligned`), C = B-aligned (the + 0.4002 @100ep record; `--legacy` reverts). Checkpoint conventions: see the + A and C ablation sections. +- Study entrypoints match the record configs (2026-10-07): `scripts/07_train.py` + (A) was already A-native all-defaults; `scripts_monai/02_monai_train.py` + (C) now trains B-aligned by default (its record config) — `--legacy` reverts + to the pre-2026-10-07 config. C checkpoints carry a `b_aligned` marker that + 03/04 read for the Z-score input; in C's B-aligned study mode pseudo rows + enter at full weight (row "w" and `--lr` unused — the same treatment + Pipeline B's study gives them; per-sample down-weighting is `--legacy` only). - 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 +- Budget knobs: `--train-epochs` (default 100), `--max-rows` (default 80% of + split_train = 2750, also B's `--nnu-cases`), `--max-eval-rows` (default 10% + of split_test = 40), `--no-tta`. +- `--skip-train` is A only: `runs/round0/best.pt` (trained 2026-09-25 on the + previous split — a few patients moved between train/test, small leakage + caveat). +- B's one-time cost scales with `--nnu-cases` (plan_and_preprocess ≈ 3.5 + s/case at 24 cases with 8 preprocess procs). +- Historical runs (10/20-epoch legacy A/B/C, nn-loss experiments, smoke run) + remain in `results/benchmark_full/` and `results/backup_20261004/`; + superseded analyses (fg-crop/CE-weight, loss parity) are omitted from this + file — see the TTA-bug section for the correction. \ No newline at end of file diff --git a/README.md b/README.md index cc91ff6..67e5565 100644 --- a/README.md +++ b/README.md @@ -63,7 +63,10 @@ logs/ run logs ## Pipeline A — In-house 3D U-Net (`scripts/`) 3-class-capable but used as binary (background / tumor) `Unet3D`, 96³ patches, -DDP, per-sample weighted loss. +DDP. Training uses A's own all-default config (BatchNorm, ReLU, per-sample +0.5CE+0.5mean-dice, AdamW 3e-4 + warmup/cosine, grad clip 5, percentile input, +legacy augmentation, 50% fg patches) — A's best benchmark config (0.5508 +@100ep; see BENCHMARKING.md). | # | Script | Purpose | |---|---|---| @@ -248,17 +251,26 @@ frame-invariant consistency filter, and evaluation protocol so results are directly comparable to Pipelines A and B. - **Network:** `monai.networks.nets.UNet` (1 in / 2 out, channels 16→128, - ~1.2 M params) — the MONAI analogue of Pipeline A's `Unet3D(base=16, depth=4)`. -- **Data:** MONAI `Dataset` + transform chain (load → min-pad → per-axis flips / - rot90 / brightness / noise → `RandCropByPosNegLabeld` 96³). Patch sampling is - foreground-aware (pos=neg=1) rather than Pipeline A's uniform random crop. -- **Loss:** per-sample weighted Dice+CE on MONAI `DiceLoss` (background+foreground - averaged, dice weight 0.5) — the same convention as Pipeline A, so pseudo rows - are down-weighted sample-by-sample (`--pseudo-weight 0.3`). -- **Schedule:** AdamW + linear warmup + cosine annealing (Pipeline A's schedule). + ~1.2 M params) — the MONAI analogue of Pipeline A's `Unet3D(base=16, depth=4)`; + the B-aligned default adds 1×1 deep-supervision aux heads. +- **Config:** the default is the B-aligned protocol — C's best benchmark config + (0.4002 @100ep, its record; see BENCHMARKING.md): InstanceNorm(affine) + + kaiming init, B's actual loss (unweighted 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, restarted + per round), per-volume Z-score input, B's augmentation, 33% foreground patch + sampling, grad clip 12. In this mode the batch-level loss ignores row "w" and + `--lr` — pseudo rows enter at full weight, the same treatment Pipeline B's + study gives them (quality is controlled by the selection gates). `--legacy` + reverts to the pre-2026-10-07 config (BatchNorm, per-sample weighted Dice+CE + so pseudo rows are down-weighted sample-by-sample (`--pseudo-weight 0.3`), + AdamW + linear warmup + cosine, legacy augmentation (per-axis flips / rot90 / + brightness / noise), 50% foreground patches, grad clip 5, percentile [0,1] + input) — C's all-default variant is its worst benchmark config (0.2083 @100ep). - **Inference:** MONAI `sliding_window_inference` (gaussian blend, 50% overlap, - 96³ windows) + the same 4-view flip TTA as Pipeline A; multi-GPU = torchrun - rank shards (like Pipeline A), DDP for training. + 96³ windows) + the same 4-view flip TTA as Pipeline A; every checkpoint stores + a `b_aligned` marker that 03/04 read to apply the matching per-volume Z-score + input; multi-GPU = torchrun rank shards (like Pipeline A), DDP for training. - **Shared with the other pipelines:** round semantics, checkpoint layout (`runs_monai/roundK/{best,final,state}.pt`), pseudo outputs (`data/pseudo_monai/roundK/`), and the gate/consistency code itself is imported @@ -267,11 +279,11 @@ directly comparable to Pipelines A and B. | # | Script | Purpose | |---|---|---| | 01 | `01_monai_build_rows.py` | Assemble a round's training manifest (labeled + accepted pseudo rows, weighted, de-duped by key) | -| 02 | `02_monai_train.py` | DDP training of the MONAI UNet for one round (torchrun), full-weight warm start, auto-resume | +| 02 | `02_monai_train.py` | DDP training of the MONAI UNet for one round (torchrun), full-weight warm start (same config mode), auto-resume | | 03 | `03_monai_pseudo_label.py` | Sliding-window + TTA prediction of the remaining pool, selection gates + consistency filter | | 04 | `04_monai_eval_test.py` | Held-out test evaluation (Dice from probability map @0.5) | | 05 | `05_monai_run_iterative.py` | Orchestrates rounds 0…K, table + plot | -| — | `monai_common.py` | Network / weighted loss / transform chain / inference + distributed helpers | +| — | `monai_common.py` | Network (legacy + B-aligned) / losses / transform chain / inference + distributed helpers | ### Running @@ -281,10 +293,13 @@ Full study from the repo root: python scripts_monai/05_monai_run_iterative.py --rounds 4 --gpus 3 ``` -Defaults: baseline 100 epochs @ 3e-4; warm-started rounds 30 epochs @ 1e-3. -Pseudo-label gates match Pipelines A/B (`--tau-pos 0.95`, `--min-cc-frac 0.2`, -`--neg-frac 0.9`, volume p2–p98). Optional flags: `--no-tta`, -`--no-neg-pseudo`, and the gate overrides. Outputs to +Defaults: B-aligned protocol (SGD 1e-2 + PolyLR, restarted per round); baseline +100 epochs, warm-started rounds 30 epochs. `--legacy` reverts to the +pre-2026-10-07 config (baseline 100 epochs @ 3e-4, warm-started rounds +30 epochs @ 1e-3, pseudo rows down-weighted). Pseudo-label gates match +Pipelines A/B (`--tau-pos 0.95`, `--min-cc-frac 0.2`, `--neg-frac 0.9`, volume +p2–p98). Optional flags: `--no-tta`, `--no-neg-pseudo`, `--legacy`, and the gate +overrides. Outputs to `results/round{k}_test_monai.json`, `results/iterative_table_monai.jsonl`, `results/iterative_dice_monai.png`; per-round checkpoints in `runs_monai/round{k}/best.pt`. @@ -297,7 +312,7 @@ python scripts_monai/01_monai_build_rows.py --train data/manifests/split_train.j --out data/pseudo_monai/round1_rows.jsonl torchrun --standalone --nproc_per_node 3 scripts_monai/02_monai_train.py \ --rows data/pseudo_monai/round1_rows.jsonl --val data/manifests/split_val.jsonl \ - --epochs 30 --lr 1e-3 --batch 3 --ckpt-dir runs_monai/round1 \ + --epochs 30 --batch 3 --ckpt-dir runs_monai/round1 \ --pretrained runs_monai/round0/best.pt torchrun --standalone --nproc_per_node 3 scripts_monai/03_monai_pseudo_label.py \ --ckpt runs_monai/round1/best.pt --pool data/manifests/unlabeled_pool.jsonl \ diff --git a/REGISTRATION_BENCHMARK_AFFINE.md b/REGISTRATION_BENCHMARK_AFFINE.md index f94efb6..27e658c 100644 --- a/REGISTRATION_BENCHMARK_AFFINE.md +++ b/REGISTRATION_BENCHMARK_AFFINE.md @@ -207,7 +207,7 @@ for the case): source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 cd /mnt/b4/xfr/git26/longitudinal -# all nine affine pipelines (results go to results/bench_registration_affine/; +# all nine affine pipelines (results go to results/bench_reg/bench_registration_affine/; # 2026-10-03 run, 6 workers, ~50 min wall dominated by impact's ~6 min phase) python scripts/bench_registration.py --affine --pipelines "antspyx ants_trsaa ants_reproa ants_quickreproa elastix impact niftyreg simpleitk synthmorph" --workers 6 --threads 24 --qa-images @@ -221,7 +221,7 @@ python scripts/bench_registration.py --affine --merge ## Visual QA images `--qa-images` writes a `checkreg`-style PNG per (pipeline, case) to -`results/bench_registration_affine/qa/__.png` (9 per +`results/bench_reg/bench_registration_affine/qa/__.png` (9 per pipeline; the incomplete-CT case is skipped) — same format as the other benchmarks: a 3×2 grid of the mid-axial / mid-coronal / mid-sagittal slices through the head center, **left column** checkerboard blend of the warped CT diff --git a/REGISTRATION_BENCHMARK_DEFORMABLE.md b/REGISTRATION_BENCHMARK_DEFORMABLE.md index 51b29cc..583b091 100644 --- a/REGISTRATION_BENCHMARK_DEFORMABLE.md +++ b/REGISTRATION_BENCHMARK_DEFORMABLE.md @@ -411,7 +411,7 @@ established empirically (validated against each tool's own resampling, NCC source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 cd /mnt/b4/xfr/git26/longitudinal -# all six deformable pipelines (results go to results/bench_registration_deformable/) +# all six deformable pipelines (results go to results/bench_reg/bench_registration_deformable/) # the 2026-10-02 main run used --workers 6 --threads 24 (~2 h 9 min wall, # dominated by simpleitk's ~111 min phase) python scripts/bench_registration.py --deformable --workers 6 --threads 24 --qa-images @@ -420,7 +420,7 @@ python scripts/bench_registration.py --deformable --workers 6 --threads 24 --qa- python scripts/bench_registration.py --deformable --pipelines simpleitk impact --n-cases 3 # the ANTs-variant follow-up (direct SyNRA / SyN + Repro / QuickRepro -# sr/br; results go to the same results/bench_registration_deformable/ dir, +# sr/br; results go to the same results/bench_reg/bench_registration_deformable/ dir, # 2026-10-03 run, 6 workers, ~28 min wall; `ants_syn` is the re-run of the # stock SyN preset that replaced the initial SyNOnly row) python scripts/bench_registration.py --deformable --pipelines "ants_synra ants_syn ants_synreprosr ants_synreprobr ants_synquickreprosr ants_synquickreprobr" --workers 6 --threads 24 --qa-images @@ -441,7 +441,7 @@ python scripts/bench_registration.py --deformable --merge ## Visual QA images `--qa-images` writes a `checkreg`-style PNG per (pipeline, case) to -`results/bench_registration_deformable/qa/__.png` +`results/bench_reg/bench_registration_deformable/qa/__.png` (9 per pipeline; the incomplete-CT case is skipped) — same format as the rigid benchmark: a 3×2 grid of the mid-axial / mid-coronal / mid-sagittal slices through the head center, **left column** checkerboard blend of the diff --git a/REGISTRATION_BENCHMARK_RIGID.md b/REGISTRATION_BENCHMARK_RIGID.md index 23685f6..8fc8bcf 100644 --- a/REGISTRATION_BENCHMARK_RIGID.md +++ b/REGISTRATION_BENCHMARK_RIGID.md @@ -9,7 +9,7 @@ MR sessions, so the inter-session head tilt can reach ~30°. - Pipelines: [`src/registration_pipelines.py`](src/registration_pipelines.py) - Benchmark harness: [`scripts/bench_registration.py`](scripts/bench_registration.py) -- Results: `results/bench_registration/` (per-pipeline JSON, `summary.json`, +- Results: `results/bench_reg/bench_registration/` (per-pipeline JSON, `summary.json`, `summary.md`, `bench_registration.png`) ## Unified contract @@ -171,7 +171,7 @@ are re-runs under the **fixed = T1c / moving = CT** convention (the shared rule); the CT-fixed-era numbers and the retired konfai-preset run are in the git history / archived JSONs. The `impact` column is the Docker rigid MIND+TS static config (the CT-fixed-era Docker run is archived in -`results/bench_registration/impact_docker_legacy.json`). +`results/bench_reg/bench_registration/impact_docker_legacy.json`). | case | antspyx | elastix | impact | niftyreg | simpleitk | synthmorph | |---|---|---|---|---|---|---| @@ -276,7 +276,7 @@ the case): ~±0.01–0.05 MI on the weak cases). IMPACT in the **Static** config re-samples spatial coordinates every iteration, but the feature maps are deterministic given the pose. The earlier Jacobian-mode Docker experiments - (archived in `results/bench_registration/impact_*.json`) resampled patches + (archived in `results/bench_reg/bench_registration/impact_*.json`) resampled patches through the network each iteration, so their per-case numbers are one sample from a distribution — e.g. the MRSegmentator-only Jacobian configs reached 0.631 (500/2000) and 0.591 (2000/5000) at ~5×/~35× the Static @@ -479,11 +479,11 @@ python scripts/bench_registration.py --pipelines synthmorph --n-cases 10 \ # for ALL pipelines, before registration and scoring; niftyreg/synthmorph # already clip by default, so this only affects antspyx/elastix/simpleitk) python scripts/bench_registration.py --n-cases 10 --workers 4 --threads 24 \ - --out-dir results/bench_reg_noclip + --out-dir results/bench_reg/bench_reg_noclip python scripts/bench_registration.py --n-cases 10 --workers 4 --threads 24 \ - --clip-ct 0 80 --out-dir results/bench_reg_clip0_80 -python scripts/bench_registration.py --merge --out-dir results/bench_reg_noclip -python scripts/bench_registration.py --merge --out-dir results/bench_reg_clip0_80 + --clip-ct 0 80 --out-dir results/bench_reg/bench_reg_clip0_80 +python scripts/bench_registration.py --merge --out-dir results/bench_reg/bench_reg_noclip +python scripts/bench_registration.py --merge --out-dir results/bench_reg/bench_reg_clip0_80 # combine the per-pipeline JSONs accumulated across envs into one table + plot python scripts/bench_registration.py --merge @@ -822,7 +822,7 @@ re-derived per toolkit. Under the shared fixed=T1c / moving=CT rule the raw - **Stochastic.** The config is **Static** mode: it re-samples spatial coordinates every iteration, but the feature maps are deterministic given the pose (see Stochasticity). The older Jacobian-mode configs - (docker-era experiments, archived in `results/bench_registration/impact_*.json`) + (docker-era experiments, archived in `results/bench_reg/bench_registration/impact_*.json`) resampled patches through the network each iteration, so their per-case numbers are one sample from a distribution. - **Models are under the data root, not the repo.** The TorchScript diff --git a/scripts/bench_registration.py b/scripts/bench_registration.py index 04365b3..0675bc3 100644 --- a/scripts/bench_registration.py +++ b/scripts/bench_registration.py @@ -624,12 +624,12 @@ def parse_args(): help="run the affine (12-DOF) variants of the pipelines " "(each toolkit's affine workflow + the antspyx " "Affine / TRSAA / Repro[a] / QuickRepro[a] presets; " - "results go to results/bench_registration_affine " - "by default)") +"results go to results/bench_reg/bench_registration_affine " + "by default)") ap.add_argument("--deformable", action="store_true", help="run the deformable variants of the pipelines " - "(each toolkit's native deformable workflow; results go " - "to results/bench_registration_deformable by default)") +"(each toolkit's native deformable workflow; results go " + "to results/bench_reg/bench_registration_deformable by default)") ap.add_argument("--pipelines", default=" ".join(RP.ALL_PIPELINES), help="space-separated subset of: " + " ".join(RP.ALL_PIPELINES)) ap.add_argument("--n-cases", type=int, default=10, help="number of M6 test cases (default 10)") @@ -642,7 +642,7 @@ def parse_args(): ap.add_argument("--simpleitk-mode", choices=["benchmark", "production"], default="benchmark", help="benchmark = single pyramid (comparable budget to other toolkits); " "production = full multi-start tilt grid (slow)") - ap.add_argument("--out-dir", default=os.path.join(ROOT, "results", "bench_registration")) + ap.add_argument("--out-dir", default=os.path.join(ROOT, "results", "bench_reg", "bench_registration")) ap.add_argument("--elastix-dir", default=None, help="elastix scratch dir (default data/m6reg/elastix)") ap.add_argument("--synthmorph-model", choices=["rigid", "affine"], default="rigid", help="SynthMorph matrix model (default rigid; deformable/joint are " @@ -692,11 +692,11 @@ def parse_args(): def main(): args = parse_args() if args.deformable: - if args.out_dir == os.path.join(ROOT, "results", "bench_registration"): - args.out_dir = os.path.join(ROOT, "results", "bench_registration_deformable") + if args.out_dir == os.path.join(ROOT, "results", "bench_reg", "bench_registration"): + args.out_dir = os.path.join(ROOT, "results", "bench_reg", "bench_registration_deformable") if args.affine: - if args.out_dir == os.path.join(ROOT, "results", "bench_registration"): - args.out_dir = os.path.join(ROOT, "results", "bench_registration_affine") + if args.out_dir == os.path.join(ROOT, "results", "bench_reg", "bench_registration"): + args.out_dir = os.path.join(ROOT, "results", "bench_reg", "bench_registration_affine") if args.merge: merge(args) return diff --git a/scripts/benchmark_pipelines.py b/scripts/benchmark_pipelines.py index 7ebf753..e7ff0ac 100644 --- a/scripts/benchmark_pipelines.py +++ b/scripts/benchmark_pipelines.py @@ -9,12 +9,24 @@ Protocol 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 + 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, 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. + 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) @@ -35,11 +47,13 @@ python scripts/benchmark_pipelines.py --pipelines A C # 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) +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 @@ -86,15 +100,57 @@ def parse_args(): "(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 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). " + 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 " @@ -102,7 +158,11 @@ def parse_args(): "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)") + 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() @@ -137,6 +197,8 @@ 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: @@ -147,7 +209,37 @@ def select_rows(args): 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] + 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 ---------------- @@ -201,18 +293,84 @@ def load_native_label(lab_path, img_itk): # ---------------- Pipeline A ---------------- -def bench_A(args, train_rows, test_rows): +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 + 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 - res = {"pipeline": "A", "backbone": "Unet3D(base=16, depth=4)", "tta": tta, - "n_train": len(train_rows), "n_eval": len(test_rows)} + 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): @@ -220,27 +378,45 @@ def bench_A(args, train_rows, test_rows): model, _ = A.load_model(ckpt, dev) res["ckpt"] = ckpt else: - model = A.build_model(base=16, device="cuda") + 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) + seed=0, persistent=False, + pos_neg=0.33 if ba else 0.5, + zscore=ba, b_aug=ba) 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]) - ckpt = os.path.join(ckpt_dir, "best.pt") + 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: + 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 @@ -255,35 +431,97 @@ def bench_A(args, train_rows, test_rows): 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) - 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) + 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(), 5.0) + torch.nn.utils.clip_grad_norm_(model.parameters(), 12.0 if ba else 5.0) opt.step() - sched.step() - loss_sum += float(losses.detach().mean()) + 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 - 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")) + 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) @@ -294,33 +532,17 @@ def bench_A(args, train_rows, test_rows): 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) + args.patch, step, tta, zscore=ba) torch.cuda.synchronize(dev) reset_peak(dev) - 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(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]] - pmaxs.append(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)) + 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.max()), 4)} + "pmax": round(float(pm), 4)} for i, pm in zip(range(len(dices)), pmaxs)] del model free_gpu() @@ -329,18 +551,67 @@ def bench_A(args, train_rows, test_rows): # ---------------- Pipeline C (MONAI) ---------------- -def bench_C(args, train_rows, test_rows): +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, make_dataloader, - predict_probs, prob_dice, load_model) + 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)", "tta": tta, - "n_train": len(train_rows), "n_eval": len(test_rows)} + 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): @@ -348,26 +619,45 @@ def bench_C(args, train_rows, test_rows): 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) + 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) - 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() - ckpt = os.path.join(ckpt_dir, "best.pt") + 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: + 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 @@ -382,31 +672,77 @@ def bench_C(args, train_rows, test_rows): 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) - loss = loss_f(logits, lab, wts) + 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(), 5.0) + torch.nn.utils.clip_grad_norm_(model.parameters(), 12.0 if ba else 5.0) opt.step() - sched.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 - 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")) + 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) @@ -414,22 +750,10 @@ def bench_C(args, train_rows, test_rows): 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) + tta=tta, sw_batch=8, zscore=ba) 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)) + 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) @@ -444,7 +768,7 @@ def bench_C(args, train_rows, test_rows): # ---------------- Pipeline B (nnU-Net) ---------------- -def bench_B(args, train_rows, test_rows): +def bench_B(args, train_rows, test_rows, val_rows=None): import nnu_common as N import SimpleITK as sitk @@ -515,7 +839,7 @@ def bench_B(args, train_rows, test_rows): raise FileNotFoundError( f"--skip-train: {N.best_ckpt()} not found (study nnU-Net results missing)") - ckpt = N.best_ckpt() + 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 @@ -526,7 +850,7 @@ def bench_B(args, train_rows, test_rows): 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") + checkpoint_name=args.nnu_ckpt) net = predictor.network if hasattr(net, "module"): net = net.module @@ -627,6 +951,108 @@ def print_table(results): print() +def _dump_partial(res, args): + """Write the in-progress res dict to .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 @@ -666,17 +1092,42 @@ def plot(results, out_png, epochs): 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.json") + 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 = select_rows(args) + 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"epochs={args.train_epochs} tta={not args.no_tta} skip_train={args.skip_train}") + 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 = [] @@ -684,7 +1135,7 @@ def main(): t0 = time.time() print(f"\n[bench] ---- pipeline {p} ----") try: - r = runners[p](args, train_rows, test_rows) + 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() @@ -696,15 +1147,25 @@ def main(): 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"), +"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: @@ -713,6 +1174,19 @@ def main(): 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__": diff --git a/scripts/run_b_100ep_stages.sh b/scripts/run_b_100ep_stages.sh new file mode 100644 index 0000000..6f32fab --- /dev/null +++ b/scripts/run_b_100ep_stages.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +# B 20->100 extension: two training stages + final best-ckpt eval. +# Trajectory points use --nnu-ckpt checkpoint_final.pth (model at the last +# trained epoch, A/C-comparable); stage 3 scores checkpoint_best.pth +# (the repo convention for B's official dice). +# bash scripts/run_b_100ep_stages.sh > logs/bench_B_100ep.log 2>&1 +set -euo pipefail + +source /opt/conda/etc/profile.d/conda.sh +conda activate longitudinal-313 + +NNU_ROOT=/home/xfr/longitudinal_nnu +BENCH_B_DIR=/home/xfr/longitudinal_benchB +export NNU_ROOT BENCH_B_DIR + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +FOLD="$NNU_ROOT/results/Dataset221_T1C_BENCH/NTUHLPLTrainer__nnUNetPlans__3d_fullres/fold_0" +GPU="${1:-0}" + +echo "=== STAGE1 20->50 start $(date -u +%FT%TZ) (gpu $GPU) ===" +cp "$FOLD/checkpoint_final.pth" "$FOLD/checkpoint_latest.pth" +python "$ROOT/scripts/benchmark_pipelines.py" --pipelines B --gpu "$GPU" \ + --resume --nnu-skip-prep --train-epochs 50 --nnu-ckpt checkpoint_final.pth \ + --out "$ROOT/results/benchmark_full/benchmark_full_B_50ep.json" + +echo "=== STAGE2 50->100 start $(date -u +%FT%TZ) (gpu $GPU) ===" +cp "$FOLD/checkpoint_final.pth" "$FOLD/checkpoint_latest.pth" +python "$ROOT/scripts/benchmark_pipelines.py" --pipelines B --gpu "$GPU" \ + --resume --nnu-skip-prep --train-epochs 100 --nnu-ckpt checkpoint_final.pth \ + --out "$ROOT/results/benchmark_full/benchmark_full_B_100ep.json" + +echo "=== STAGE3 best-ckpt eval start $(date -u +%FT%TZ) (gpu $GPU) ===" +python "$ROOT/scripts/benchmark_pipelines.py" --pipelines B --gpu "$GPU" \ + --skip-train --nnu-dsid 221 --nnu-ckpt checkpoint_best.pth \ + --out "$ROOT/results/benchmark_full/benchmark_full_B_100ep_best.json" + +echo "=== ALL B STAGES DONE $(date -u +%FT%TZ) ===" \ No newline at end of file diff --git a/scripts/run_c_bloss_100ep_after_b.sh b/scripts/run_c_bloss_100ep_after_b.sh new file mode 100644 index 0000000..37a2dca --- /dev/null +++ b/scripts/run_c_bloss_100ep_after_b.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Wait for B's 100-ep stages (stage 2 train + stage 3 best-ckpt inference) to +# release GPU 0, then run the C b-loss-only 100-ep ablation. +source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313 +cd /mnt/b4/xfr/git26/longitudinal +echo "waiting for GPU0 (B 100ep stages) at $(date -u)" +while true; do + mem=$(nvidia-smi --id=0 --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | tr -d ' ') + util=$(nvidia-smi --id=0 --query-gpu=utilization.gpu --format=csv,noheader,nounits 2>/dev/null | tr -d ' ') + if [ -n "$mem" ] && [ "$mem" -lt 1000 ] && [ -n "$util" ] && [ "$util" -lt 5 ]; then + echo "GPU0 free (mem=${mem}MiB util=${util}%) at $(date -u); launching C b-loss-only 100ep" + break + fi + sleep 120 +done +python scripts/benchmark_pipelines.py --pipelines C --b-loss-only --gpu 0 --train-epochs 100 \ + --eval-at "5 10 20 50" --workers 4 \ + --out results/benchmark_full/benchmark_full_C_bloss_100ep.json 2>&1 | tee logs/bench_C_bloss_100ep.log +echo "C b-loss-only 100ep finished at $(date -u)" \ No newline at end of file diff --git a/scripts_monai/02_monai_train.py b/scripts_monai/02_monai_train.py index 3e84b26..222bd7c 100644 --- a/scripts_monai/02_monai_train.py +++ b/scripts_monai/02_monai_train.py @@ -2,17 +2,28 @@ torchrun --standalone --nproc_per_node 3 scripts_monai/02_monai_train.py \ --rows data/pseudo_monai/round0_rows.jsonl --val data/manifests/split_val.jsonl \ - --epochs 100 --lr 3e-4 --batch 3 --ckpt-dir runs_monai/round0 \ + --epochs 100 --batch 3 --ckpt-dir runs_monai/round0 \ [--pretrained runs_monai/round0/best.pt] -- backbone: monai.networks.nets.UNet (monai_common.build_model) -- data: MONAI Dataset + transform chain (load → pad → flip/rotate90/intensity → - foreground-aware 96³ crop → typed tensors); row "w" carries the per-sample - loss weight (pseudo rows down-weighted) -- loss: per-sample Dice+CE, weighted batch mean (Pipeline A's convention) -- schedule: AdamW + linear warmup + cosine (Pipeline A's schedule) +Config modes (see BENCHMARKING.md, "B-aligned A/C" + "C ablation" sections): +- DEFAULT: the B-aligned protocol — C's best benchmark config (0.4002 @100ep, + its record). monai UNet(16→128) with InstanceNorm(affine) + kaiming init, + B's ACTUAL loss (b_dc_and_ce: 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 this round's + epoch budget — restarted per round, the staged convention), per-volume + Z-score input, B's augmentation, 33% foreground patch sampling, grad clip 12. + Batch-level loss: row "w" is NOT used (pseudo rows enter at full weight, + the same treatment Pipeline B's study gives them — quality is controlled by + the selection gates), and --lr is ignored (protocol-fixed 1e-2). +- --legacy: the pre-2026-10-07 config (BatchNorm, per-sample weighted Dice+CE + with row "w" down-weighting, AdamW + linear warmup + cosine, legacy + augmentation, 50% fg patches, grad clip 5, percentile [0,1] input) — + C's worst benchmark config (0.2083 @100ep all-default). + - checkpoints in --ckpt-dir: best.pt (max val dice), final.pt, state.pt - (auto-resume after a crash, as in Pipeline A) + (auto-resume after a crash); every save carries a "b_aligned" marker that + 03/04 read for the matching Z-score input protocol """ import argparse import os @@ -22,9 +33,12 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import numpy as np import torch +import torch.nn.functional as F from src.common import load_jsonl, read_nii_arr +from src.losses import b_dc_and_ce, deep_sup_loss from monai_common import (build_model, WeightedDiceCELoss, make_dataloader, - predict_probs, prob_dice, rank_info, init_dist, barrier, destroy_dist) + predict_probs, prob_dice, rank_info, init_dist, barrier, + destroy_dist, ckpt_b_aligned) def main(): @@ -32,16 +46,23 @@ def main(): ap.add_argument("--rows", required=True, help="jsonl of {key,pimg,plabel,w}") ap.add_argument("--val", required=True, help="held-out val split jsonl ({pimg,plabel})") ap.add_argument("--epochs", type=int, default=100) - ap.add_argument("--lr", type=float, default=3e-4) + ap.add_argument("--lr", type=float, default=3e-4, + help="legacy mode only (AdamW base lr); ignored in B-aligned mode " + "(protocol-fixed SGD 1e-2 + PolyLR)") ap.add_argument("--batch", type=int, default=3, help="per-GPU batch") ap.add_argument("--patch", type=int, default=96) ap.add_argument("--workers", type=int, default=4) ap.add_argument("--ckpt-dir", required=True) - ap.add_argument("--pretrained", default=None, help="full-weight warm-start checkpoint") + ap.add_argument("--pretrained", default=None, help="full-weight warm-start checkpoint (same config mode)") + ap.add_argument("--legacy", action="store_true", + help="pre-2026-10-07 C config (BatchNorm, per-sample weighted Dice+CE, " + "AdamW + warmup/cosine, legacy augmentation, 50%% fg, clip 5); the " + "default is the B-aligned protocol (C's record benchmark config)") ap.add_argument("--val-every", type=int, default=10) ap.add_argument("--val-limit", type=int, default=60) ap.add_argument("--sw-batch", type=int, default=8, help="sliding-window batch at inference") args = ap.parse_args() + ba = not args.legacy rank, world, local_rank = rank_info() init_dist() @@ -53,47 +74,68 @@ def main(): rows = load_jsonl(args.rows) val_rows = load_jsonl(args.val)[:args.val_limit] - dl = make_dataloader(rows, win=args.patch, batch=args.batch, workers=args.workers) + dl = make_dataloader(rows, win=args.patch, batch=args.batch, workers=args.workers, + b_aligned=ba) steps_per_epoch = max(len(dl), 1) - model = build_model(device=device) + model = build_model(device, norm="instance" if ba else "batch", b_aligned=ba) if args.pretrained: sd0 = torch.load(args.pretrained, map_location=device, weights_only=True) + if ckpt_b_aligned(sd0) != ba: + raise SystemExit(f"--pretrained {args.pretrained} is " + f"{'B-aligned' if ckpt_b_aligned(sd0) else 'legacy'}-config and cannot " + f"warm-start the {'B-aligned' if ba else 'legacy'} model (norm layers + " + f"aux heads differ); retrain or drop --pretrained") model.load_state_dict(sd0.get("model", sd0)) if rank == 0: print(f"[monai:train:rank0] warm start (all weights) from {args.pretrained}", flush=True) ddp = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank]) if world > 1 else model total_steps = steps_per_epoch * args.epochs - opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) - warmup = min(300, max(10, total_steps // 10)) - base_sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1), - eta_min=args.lr * 0.05) - sched = torch.optim.lr_scheduler.SequentialLR( - opt, [torch.optim.lr_scheduler.LinearLR(opt, start_factor=0.1, total_iters=warmup), base_sched], - milestones=[warmup]) - loss_f = WeightedDiceCELoss() + if ba: + opt = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.99, nesterov=True, + weight_decay=3e-5) + sched = None # PolyLR is stateless: lr recomputed per epoch over this round's budget + else: + opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) + warmup = min(300, max(10, total_steps // 10)) + base_sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1), + eta_min=args.lr * 0.05) + sched = torch.optim.lr_scheduler.SequentialLR( + opt, [torch.optim.lr_scheduler.LinearLR(opt, start_factor=0.1, total_iters=warmup), base_sched], + milestones=[warmup]) + loss_f = WeightedDiceCELoss() if not ba else None best, best_epoch, start_epoch = -1.0, -1, 0 if rank == 0: os.makedirs(args.ckpt_dir, exist_ok=True) + print(f"[monai:train:rank0] config={'B-aligned' if ba else 'legacy'} " + f"rows={len(rows)} val={len(val_rows)} epochs={args.epochs} " + f"steps/epoch={steps_per_epoch} world={world}", flush=True) + if ba: + print("[monai:train:rank0] B-aligned: SGD 1e-2 + PolyLR(0.9), b_dc_and_ce + deep sup " + "(4/7, 2/7, 1/7), Z-score input, B aug, 33% fg, clip 12; --lr and row 'w' unused", + flush=True) state_f = os.path.join(args.ckpt_dir, "state.pt") if os.path.exists(state_f): st = torch.load(state_f, map_location="cpu", weights_only=True) + if st.get("b_aligned", False) != ba: + raise SystemExit(f"{state_f} is {'B-aligned' if st.get('b_aligned') else 'legacy'}-config, " + f"this run is {'B-aligned' if ba else 'legacy'}; remove state.pt or match --legacy") model.load_state_dict(st["model"]) opt.load_state_dict(st["opt"]) - sched.load_state_dict(st["sched"]) + if sched is not None: + sched.load_state_dict(st["sched"]) best, best_epoch, start_epoch = st["best"], st["best_epoch"], st["epoch"] if rank == 0: print(f"[monai:train] auto-resuming from state.pt @epoch {start_epoch} (best={best:.4f})", flush=True) - if rank == 0: - print(f"[monai:train] rows={len(rows)} val={len(val_rows)} epochs={args.epochs} " - f"steps/epoch={steps_per_epoch} world={world}", flush=True) barrier() def save_state(): tmp = state_f + ".tmp" - torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "sched": sched.state_dict(), + torch.save({"model": model.state_dict(), "opt": opt.state_dict(), + "sched": sched.state_dict() if sched is not None else None, + "b_aligned": ba, "best": best, "best_epoch": best_epoch, "epoch": epoch + 1}, tmp) os.replace(tmp, state_f) @@ -101,6 +143,11 @@ def main(): if hasattr(dl.sampler, "set_epoch"): dl.sampler.set_epoch(epoch) model.train() + if ba: + # PolyLR (B): lr = 1e-2 * (1 - ep/total)^0.9 on the GLOBAL epoch index + lr0 = 1e-2 * (1 - epoch / args.epochs) ** 0.9 + for pg in opt.param_groups: + pg["lr"] = lr0 run_loss, run_n = 0.0, 0 for batch in dl: img = batch["pimg"].to(device, non_blocking=True) @@ -108,12 +155,22 @@ def main(): wts = batch["w"].to(device, non_blocking=True) with torch.autocast("cuda", dtype=torch.bfloat16): logits = ddp(img) - loss = loss_f(logits, lab, wts) + if ba: + l2 = F.interpolate(lab.unsqueeze(1).float(), scale_factor=0.5, + mode="nearest").squeeze(1).long() + l4 = F.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) + if not torch.isfinite(loss): + loss = torch.zeros_like(loss) # zero-gradient fallback: keeps DDP collectives in sync + else: + 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) + torch.nn.utils.clip_grad_norm_(model.parameters(), 12.0 if ba else 5.0) opt.step() - sched.step() + if sched is not None: + sched.step() run_loss += float(loss.detach()) run_n += 1 if rank == 0: @@ -122,31 +179,39 @@ def main(): if rank == 0 and (epoch + 1) % max(1, args.val_every // 2) == 0: save_state() if val_rows and rank == 0 and (epoch + 1) % args.val_every == 0: + # eval mode: UNetBAligned's deep-sup forward returns a tuple in train mode + model.eval() dv, n_ok = 0.0, 0 for r in val_rows: try: vol = read_nii_arr(r["pimg"]).astype("float32") labv = read_nii_arr(r["plabel"]).astype("uint8") - p = predict_probs(model, vol, device, args.patch, tta=True, sw_batch=args.sw_batch) + p = predict_probs(model, vol, device, args.patch, tta=True, + sw_batch=args.sw_batch, zscore=ba) if p.shape == labv.shape: dv += prob_dice(p, labv) n_ok += 1 except Exception as e: # noqa print(" val err", r.get("key"), repr(e)) + model.train() dv = dv / max(n_ok, 1) print(f" [val] epoch {epoch+1} dice={dv:.4f}", flush=True) if dv > best: - best, best_epoch = dv, epoch + 1 - torch.save({"model": model.state_dict(), "epoch": epoch + 1, "val_dice": best}, + best = dv + best_epoch = epoch + 1 + torch.save({"model": model.state_dict(), "epoch": epoch + 1, "val_dice": best, + "b_aligned": ba}, os.path.join(args.ckpt_dir, "best.pt")) save_state() barrier() if rank == 0: - torch.save({"model": model.state_dict(), "epoch": args.epochs, "val_dice": best}, + torch.save({"model": model.state_dict(), "epoch": args.epochs, "val_dice": best, + "b_aligned": ba}, os.path.join(args.ckpt_dir, "final.pt")) best_f = os.path.join(args.ckpt_dir, "best.pt") if not os.path.exists(best_f): - torch.save({"model": model.state_dict(), "epoch": args.epochs, "val_dice": 0.0}, best_f) + torch.save({"model": model.state_dict(), "epoch": args.epochs, "val_dice": 0.0, + "b_aligned": ba}, best_f) print("[monai:train] no val run; best.pt = final weights", flush=True) print(f"[monai:train] done best_val_dice={best:.4f}@{best_epoch}", flush=True) destroy_dist() diff --git a/scripts_monai/03_monai_pseudo_label.py b/scripts_monai/03_monai_pseudo_label.py index 5bc02ea..3f0bee0 100644 --- a/scripts_monai/03_monai_pseudo_label.py +++ b/scripts_monai/03_monai_pseudo_label.py @@ -33,7 +33,7 @@ import SimpleITK as sitk from src.common import load_jsonl, save_jsonl, read_nii_arr, write_arr, d from monai_common import (rank_info, init_dist, barrier, destroy_dist, load_model, predict_probs, pos_mask, neg_frac_bg, consistency_filter, - load_voxel_stats) + load_voxel_stats, ckpt_b_aligned) def main(): @@ -84,7 +84,10 @@ def main(): f"done={len(done)} to_predict={len(rows)} shard={len(shard)} " f"vol_range=[{vol_lo:.0f},{vol_hi:.0f}]mm3", flush=True) - model, _ = load_model(args.ckpt, device) + model, sd = load_model(args.ckpt, device) + ba = ckpt_b_aligned(sd) # B-aligned models need the per-volume Z-score input + if rank == 0: + print(f"[monai:pseudo:rank0] ckpt config={'B-aligned' if ba else 'legacy'}", flush=True) part = os.path.join(out, f"part{rank}.jsonl") buf, n_err = [], 0 for i, r in enumerate(shard, 1): @@ -96,7 +99,7 @@ def main(): pimg_itk = sitk.ReadImage(r["pimg"]) vol = read_nii_arr(r["pimg"]).astype("float32") pt = predict_probs(model, vol, device, args.win, args.overlap, - tta=not args.no_tta, sw_batch=args.sw_batch) + tta=not args.no_tta, sw_batch=args.sw_batch, zscore=ba) if pt.shape != vol.shape: raise ValueError(f"prob shape {pt.shape} != image shape {vol.shape} for {key}") entry["maxp"] = round(float(pt.max()), 4) @@ -157,7 +160,7 @@ def main(): "tau_pos": args.tau_pos, "tau_neg": args.tau_neg, "neg_frac": args.neg_frac, "vol_range": [vol_lo, vol_hi], "max_rel_dist_mm": args.max_rel_dist, "vol_ratio": args.vol_ratio, "win": args.win, "overlap": args.overlap, - "tta": not args.no_tta, "ckpt": args.ckpt, + "tta": not args.no_tta, "ckpt": args.ckpt, "b_aligned": ba, } with open(os.path.join(out, "summary.json"), "w") as f: json.dump(summ, f, indent=1) diff --git a/scripts_monai/04_monai_eval_test.py b/scripts_monai/04_monai_eval_test.py index 79e04ef..874a9e5 100644 --- a/scripts_monai/04_monai_eval_test.py +++ b/scripts_monai/04_monai_eval_test.py @@ -25,7 +25,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import torch from src.common import load_jsonl, save_jsonl, read_nii_arr from monai_common import (rank_info, init_dist, barrier, destroy_dist, load_model, - predict_probs, prob_dice) + predict_probs, prob_dice, ckpt_b_aligned) def main(): @@ -55,7 +55,8 @@ def main(): print(f"[monai:eval:rank0] n_test={len(rows)} done={len(done)} " f"shard(world={world})={len(shard)}", flush=True) - model, _ = load_model(args.ckpt, device) + model, sd = load_model(args.ckpt, device) + ba = ckpt_b_aligned(sd) # B-aligned models need the per-volume Z-score input with open(outp, "a") as f: for i, r in enumerate(shard, 1): key = r["key"] @@ -64,7 +65,7 @@ def main(): vol = read_nii_arr(r["pimg"]).astype("float32") lab = read_nii_arr(r.get("plabel") or r.get("label")).astype("uint8") p = predict_probs(model, vol, device, args.win, args.overlap, - tta=not args.no_tta, sw_batch=args.sw_batch) + tta=not args.no_tta, sw_batch=args.sw_batch, zscore=ba) if p.shape == lab.shape: entry["dice"] = round(prob_dice(p, lab), 4) else: @@ -89,7 +90,7 @@ def main(): dice_ok = [e["dice"] for e in per if e["dice"] is not None] res = {"ckpt": args.ckpt, "n": len(per), "n_ok": len(dice_ok), "dice": float(sum(dice_ok) / len(dice_ok)) if dice_ok else 0.0, - "tta": not args.no_tta} + "tta": not args.no_tta, "b_aligned": ba} os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) with open(args.out, "w") as f: json.dump(res, f, indent=1) diff --git a/scripts_monai/05_monai_run_iterative.py b/scripts_monai/05_monai_run_iterative.py index a628e25..f8171e5 100644 --- a/scripts_monai/05_monai_run_iterative.py +++ b/scripts_monai/05_monai_run_iterative.py @@ -5,11 +5,17 @@ Round 0: train the MONAI UNet on the labeled patient-level train split (internal validation = held-out patient-level val split, subject-disjoint). Round k: dataset grows with accepted pseudo-labels (pos + neg) from rounds 1..k (de-duplicated by key, pseudo rows weighted); warm-start training - (full weights) from round k-1's best checkpoint at a lower LR; then - pseudo-label the remaining unlabeled pool (gates identical to - Pipelines A/B) and evaluate the model on the held-out test split. + (full weights) from round k-1's best checkpoint; then pseudo-label + the remaining unlabeled pool (gates identical to Pipelines A/B) and + evaluate the model on the held-out test split. -Usage: python scripts_monai/05_monai_run_iterative.py [--rounds 4] [--gpus 3] +Config: rounds train with the B-aligned protocol (C's record benchmark config, + SGD 1e-2 + PolyLR restarted per round, b_dc_and_ce + deep supervision; + pseudo rows enter at full weight, as in Pipeline B's study) unless + --legacy (pre-2026-10-07 config: AdamW at --base-lr/--round-lr + + per-sample weighted Dice+CE with pseudo-row down-weighting). + +Usage: python scripts_monai/05_monai_run_iterative.py [--rounds 4] [--gpus 3] [--legacy] """ import os import sys @@ -19,7 +25,7 @@ import argparse import subprocess sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from src.common import ROOT, d, load_jsonl, save_jsonl +from src.common import ROOT, DATA, d, load_jsonl, save_jsonl S = os.path.dirname(os.path.abspath(__file__)) # code lives next to this file, not under ROOT @@ -28,6 +34,11 @@ def main(): ap = argparse.ArgumentParser() ap.add_argument("--rounds", type=int, default=4) ap.add_argument("--gpus", type=int, default=3) + ap.add_argument("--legacy", action="store_true", + help="train all rounds with C's pre-2026-10-07 config (BatchNorm, " + "per-sample weighted Dice+CE, AdamW + warmup/cosine); the default is " + "the B-aligned protocol (C's record benchmark config, see " + "BENCHMARKING.md). --base-lr/--round-lr only apply in legacy mode") ap.add_argument("--base-epochs", type=int, default=100) ap.add_argument("--base-lr", type=float, default=3e-4) ap.add_argument("--round-epochs", type=int, default=30) @@ -47,7 +58,7 @@ def main(): ap.add_argument("--vol-ratio", type=float, default=10.0) args = ap.parse_args() - man = os.path.join(ROOT, "data/manifests") + man = os.path.join(DATA, "manifests") train_f = os.path.join(man, "split_train.jsonl") val_f = os.path.join(man, "split_val.jsonl") test_f = os.path.join(man, "split_test.jsonl") @@ -98,6 +109,8 @@ def main(): f"--val-every {val_every} --val-limit {args.val_limit}") if warmstart: extra += f" --pretrained {warmstart}" + if args.legacy: + extra += " --legacy" trrun("02_monai_train.py", extra, os.path.join(logdir, f"train_r{round}.log")) def predict_pool(round, ckpt, added=None): diff --git a/scripts_monai/monai_common.py b/scripts_monai/monai_common.py index 77f8614..5abc74a 100644 --- a/scripts_monai/monai_common.py +++ b/scripts_monai/monai_common.py @@ -30,6 +30,7 @@ import numpy as np import torch import torch.nn as nn +from monai.networks.nets import UNet # module is monai-specific; safe to import eagerly from src.common import dice # noqa: E402 # Selection gates / consistency filter shared with Pipeline B (pure numpy/scipy) @@ -40,8 +41,10 @@ WIN = 96 # ---------------- network / loss ---------------- -def build_model(device): +def build_model(device, norm="batch", b_aligned=False): from monai.networks.nets import UNet + if b_aligned: + return UNetBAligned(device, norm=norm).to(device) net = UNet( spatial_dims=3, in_channels=1, @@ -50,11 +53,49 @@ def build_model(device): strides=(2, 2, 2), num_res_units=2, act=("LEAKYRELU", {"inplace": True, "negative_slope": 0.01}), - norm="batch", + norm=("instance", {"affine": True, "eps": 1e-5}) if norm == "instance" else "batch", ) return net.to(device) +class UNetBAligned(UNet): + """Pipeline C's MONAI UNet + Pipeline B's deep supervision. + + Adds 1x1 conv aux heads at the 1/2- and 1/4-resolution decoder outputs + (the 1/8 bottleneck is dropped, mirroring nnU-Net's zero-weight lowest + scale). In training mode forward returns (final, aux_half, aux_quarter); + in eval mode it returns just the final logits so MONAI sliding-window + inference works unchanged. + """ + + def __init__(self, device, norm="instance"): + super().__init__( + spatial_dims=3, in_channels=1, out_channels=2, + channels=(16, 32, 64, 128), strides=(2, 2, 2), num_res_units=2, + act=("LEAKYRELU", {"inplace": True, "negative_slope": 0.01}), + norm=("instance", {"affine": True, "eps": 1e-5}) if norm == "instance" else "batch", + ) + # locate the up-paths: model[1].submodule = [down2, SkipConn(S3), up2] + s2 = self.model[1].submodule + self._up_half = s2[2] # 16ch @ 1/2 + s3 = s2[1].submodule # [down3, SkipConn(bottom), up3] + self._up_quarter = s3[2] # 32ch @ 1/4 + self.aux_half = nn.Conv3d(16, 2, 1) + self.aux_quarter = nn.Conv3d(32, 2, 1) + self._aux = {} + self._up_half.register_forward_hook(lambda m, i, o: self._aux.__setitem__("half", o)) + self._up_quarter.register_forward_hook(lambda m, i, o: self._aux.__setitem__("quarter", o)) + from src.unet3d import init_b + self.apply(init_b) + + def forward(self, x): + self._aux = {} + out = super().forward(x) + if not self.training: + return out + return out, self.aux_half(self._aux["half"]), self.aux_quarter(self._aux["quarter"]) + + class WeightedDiceCELoss(nn.Module): """Per-sample (soft-dice + CE), weighted batch mean. @@ -88,18 +129,108 @@ class WeightedDiceCELoss(nn.Module): # ---------------- data ---------------- -def train_transform(win=WIN): +class NNDCeLoss(nn.Module): + """nnU-Net-style DC_and_CE (same formulation as src.losses.nn_dc_and_ce): + unweighted CE + batch-level soft dice over both classes, smooth=1.0, + loss = CE - mean(dice). Used for the A/C-vs-B loss-parity runs.""" + + def forward(self, logits, label): + logits = logits.float() + ce = nn.functional.cross_entropy(logits, label) + p = nn.functional.softmax(logits, dim=1) + g = nn.functional.one_hot(label, p.size(1)).permute(0, -1, *range(1, label.ndim)).float() + axes = (0, *range(2, p.dim())) + tp = (p * g).sum(axes) + fp = (p * (1.0 - g)).sum(axes) + fn = ((1.0 - p) * g).sum(axes) + dc = (2.0 * tp + 1.0) / (2.0 * tp + fp + fn + 1.0) + loss = ce - dc.mean() + if not torch.isfinite(loss): + return torch.zeros_like(loss) # zero-gradient fallback (DDP parity) + return loss + + +class _WorkerRNG: + """Per-dataloader-worker numpy RandomState (forked workers reseed independently).""" + + @staticmethod + def get(seed=0): + info = torch.utils.data.get_worker_info() + wid = info.id if info is not None else 0 + return np.random.RandomState(seed + 104729 * wid) + + +class ZScoreAll: + """Per-volume Z-score over all voxels (Pipeline B's ZScoreNormalization, no + mask). Applied to channel-first pimg.""" + + def __call__(self, data): + from src.augment_b import zscore_volume + d = dict(data) + img = d["pimg"] + v = img[0] if img.ndim == 4 else img + v = zscore_volume(np.asarray(v, dtype=np.float32)) + d["pimg"] = v[np.newaxis] if img.ndim == 4 else v + return d + + +class BAugmentB: + """Pipeline B's full augmentation (src.augment_b), applied to the sampled + 96^3 patch (B's SpatialTransform deforms the sampled patch about its + center and crops back to the final size).""" + + def __call__(self, data): + from src.augment_b import augment_b + d = dict(data) + img = np.asarray(d["pimg"], dtype=np.float32) + raw_lab = d.get("plabel") + lab = None if raw_lab is None else np.asarray(raw_lab) + lab_4d = lab is not None and lab.ndim == 4 + if lab_4d: + lab = lab[0] + out_img, out_lab = augment_b(img, lab, _WorkerRNG.get()) + d["pimg"] = out_img if img.ndim == 4 else out_img[np.newaxis] + if out_lab is not None: + d["plabel"] = out_lab[np.newaxis] if lab_4d else out_lab + return d + + +def train_transform(win=WIN, b_aligned=False, no_aug=False): """Channel-first MONAI chain: load → min-pad → augment → pos/neg crop → typed tensors. - Mirrors Pipeline A's augmentation set (per-axis flip p=0.5, rot90 in the - (y,x) plane p=0.1, brightness 1±0.1 p=0.3, Gaussian noise σ=0.01 p=0.4). - Patch sampling uses MONAI's foreground-aware RandCropByPosNegLabeld - (pos=neg=1) instead of Pipeline A's uniform random crop. + Legacy (default): mirrors Pipeline A's augmentation set (per-axis flip p=0.5, + rot90 in the (y,x) plane p=0.1, brightness 1±0.1 p=0.3, Gaussian noise + σ=0.01 p=0.4). Patch sampling uses MONAI's foreground-aware + RandCropByPosNegLabeld (pos=neg=1). + + b_aligned: Pipeline B's data protocol — per-volume Z-score, B's full + augmentation (rot/zoom/noise/blur/brightness/contrast/lowres/gamma/mirror), + and 33% foreground patch sampling (pos=neg ratio 0.5:1, like B's + oversample_foreground_percent=0.33). """ from monai.transforms import ( Compose, LoadImaged, EnsureChannelFirstd, EnsureTyped, SpatialPadd, RandFlipd, RandRotate90d, RandScaleIntensityd, RandGaussianNoised, - RandCropByPosNegLabeld) + RandCropByPosNegLabeld, Compose as _C) + if b_aligned: + crop_pos, crop_neg = 0.5, 1.0 + # B's SpatialTransform augments the SAMPLED patch (rot/zoom about its + # center, then crop back to final size) — not the whole volume; that + # is also what Pipeline A's b_aug does, and whole-volume deformation + # is ~30x slower per item. + steps = [ + LoadImaged(keys=["pimg", "plabel"]), + EnsureChannelFirstd(keys=["pimg", "plabel"]), + SpatialPadd(keys=["pimg", "plabel"], spatial_size=(win, win, win)), + ZScoreAll(), + RandCropByPosNegLabeld(keys=["pimg", "plabel"], label_key="plabel", + spatial_size=(win, win, win), pos=crop_pos, neg=crop_neg, + num_samples=1), + ] + if not no_aug: + steps.append(BAugmentB()) + steps.append(EnsureTyped(keys=["pimg", "plabel"], dtype=[torch.float32, torch.long])) + return Compose(steps) return Compose([ LoadImaged(keys=["pimg", "plabel"]), EnsureChannelFirstd(keys=["pimg", "plabel"]), @@ -130,9 +261,9 @@ def _collate(batch): return torch.utils.data.default_collate(items) -def make_dataloader(rows, win=WIN, batch=3, workers=4): +def make_dataloader(rows, win=WIN, batch=3, workers=4, b_aligned=False): from monai.data import Dataset - ds = Dataset(data=rows_for_dataset(rows), transform=train_transform(win)) + ds = Dataset(data=rows_for_dataset(rows), transform=train_transform(win, b_aligned=b_aligned)) return torch.utils.data.DataLoader( ds, batch_size=batch, shuffle=True, num_workers=workers, collate_fn=_collate, drop_last=len(ds) > batch, pin_memory=True, @@ -142,15 +273,21 @@ def make_dataloader(rows, win=WIN, batch=3, workers=4): # ---------------- inference ---------------- @torch.no_grad() -def predict_probs(model, vol, device, win=WIN, overlap=0.5, tta=True, sw_batch=8): +def predict_probs(model, vol, device, win=WIN, overlap=0.5, tta=True, sw_batch=8, + zscore=False): """Tumor probability map (C, z, y, x → (z, y, x) numpy) for one volume. MONAI sliding_window_inference (gaussian blend, `overlap`), plus the same - 4-view flip TTA (identity + 3 axis flips) as Pipeline A. + 4-view flip TTA (identity + 3 axis flips) as Pipeline A. zscore=True matches + a B-aligned model's input (per-volume Z-score, no [0,1.5] clip). """ from monai.inferers import sliding_window_inference a = np.nan_to_num(np.asarray(vol, dtype=np.float32), nan=0.0, posinf=1.5, neginf=0.0) - a = np.clip(a, 0.0, 1.5) + if zscore: + from src.augment_b import zscore_volume + a = zscore_volume(a) + else: + a = np.clip(a, 0.0, 1.5) t = torch.from_numpy(a).unsqueeze(0).unsqueeze(0).to(device) def run(vt): @@ -175,9 +312,22 @@ def prob_dice(probs, lab, thr=0.5): return dice((probs >= thr).astype("uint8"), (lab > 0).astype("uint8")) +def ckpt_b_aligned(sd): + """Whether a checkpoint was trained with the B-aligned config (instance norm + + deep-sup aux heads). New checkpoints carry an explicit 'b_aligned' marker; + pre-marker ones fall back to the aux_head state-dict keys.""" + if not isinstance(sd, dict): + return False + if "b_aligned" in sd: + return bool(sd["b_aligned"]) + inner = sd.get("model", sd) + return isinstance(inner, dict) and any(k.startswith("aux_") for k in inner) + + def load_model(ckpt, device): sd = torch.load(ckpt, map_location=device, weights_only=True) - net = build_model(device) + ba = ckpt_b_aligned(sd) + net = build_model(device, norm="instance" if ba else "batch", b_aligned=ba) net.load_state_dict(sd.get("model", sd)) net.eval() return net, sd diff --git a/src/augment_b.py b/src/augment_b.py new file mode 100644 index 0000000..0d773df --- /dev/null +++ b/src/augment_b.py @@ -0,0 +1,149 @@ +"""Pipeline B's (nnU-Net v2 3d_fullres) training augmentation, as a numpy/scipy +implementation shared by Pipelines A and C for the B-aligned runs. + +Semantics follow the nnU-Net default 3D pipeline as resolved for Dataset221 +(dump of B's training transforms, fold_0/debug.json): + + SpatialTransform: + - rotation p=0.2 per axis, angle U(-pi/6, +pi/6), order-1 image / nn label + - scaling p=0.2, factor U(0.7, 1.4), synchronized across axes, trilinear + - elastic p=0 (dataset-specific, none for this data) + GaussianNoise p=0.1 variance U(0, 0.1) + GaussianBlur p=0.2 sigma U(0.5, 1.0) per axis (x p_per_channel 0.5) + MultiplicativeBrightness p=0.15 multiplier U(0.75, 1.25) + Contrast p=0.15 factor U(0.75, 1.25), preserve_range + SimulateLowResolution p=0.25 zoom U(0.5, 1) (down order1, up order0) + Gamma (inverted) p=0.10 gamma BGContrast(0.7, 1.5), p_retain_stats=1 + Gamma p=0.30 gamma BGContrast(0.7, 1.5), p_retain_stats=1 + Mirror 3 axes, p=0.5 each +""" +import numpy as np +from scipy import ndimage + + +def _bg_contrast_sample(rng, lo, hi): + """batchgenerators BGContrast sampling: if the range straddles 1, half the + samples come from [lo, 1), half from [1, hi].""" + if lo < 1 < hi: + if rng.random() < 0.5: + return float(rng.uniform(lo, 1.0)) + return float(rng.uniform(1.0, hi)) + return float(rng.uniform(lo, hi)) + + +def _center_crop_or_pad(a, shape, order=0, cval=0.0): + """a may be larger (crop) or smaller (pad) than shape; keep the center.""" + out = np.zeros(shape, dtype=a.dtype) + sl_in, sl_out = [], [] + for i, (sa, so) in enumerate(zip(a.shape, shape)): + if sa >= so: + start = (sa - so) // 2 + sl_in.append(slice(start, start + so)) + sl_out.append(slice(0, so)) + else: + start = (so - sa) // 2 + sl_in.append(slice(0, sa)) + sl_out.append(slice(start, start + sa)) + out[tuple(sl_out)] = a[tuple(sl_in)] + return out + + +def augment_b(img, label, rng): + """img: (1, z, y, x) float32 patch, label: (z, y, x) int (None allowed). + Returns augmented (img, label). `rng` is a numpy RandomState/Generator.""" + x = np.asarray(img[0], dtype=np.float32) + lab = None if label is None else np.asarray(label, dtype=np.int64) + + # ---- rotation (per axis, p=0.2, +-30 deg) ---- + # scipy rotate axis pairs: z-axis -> (1,2), y-axis -> (0,2), x-axis -> (0,1) + for axis_pair in ((1, 2), (0, 2), (0, 1)): + if rng.random_sample() < 0.2: + ang = float(rng.uniform(-np.pi / 6, np.pi / 6)) + x = ndimage.rotate(x, ang, axes=axis_pair, order=1, reshape=False, + mode='constant', cval=0.0) + if lab is not None: + lab = ndimage.rotate(lab, ang, axes=axis_pair, order=0, reshape=False, + mode='constant', cval=0) + + # ---- scaling (p=0.2, synchronized 0.7-1.4) ---- + if rng.random_sample() < 0.2: + f = float(rng.uniform(0.7, 1.4)) + shape = x.shape + xz = ndimage.zoom(x, f, order=1, mode='constant', cval=0.0) + x = _center_crop_or_pad(xz, shape) + if lab is not None: + lz = ndimage.zoom(lab, f, order=0, mode='constant', cval=0) + lab = _center_crop_or_pad(lz, shape) + + # ---- gaussian noise (p=0.1, variance U(0,0.1)) ---- + if rng.random_sample() < 0.1: + var = float(rng.uniform(0.0, 0.1)) + x = x + (np.sqrt(var) * rng.standard_normal(x.shape)).astype(np.float32) + + # ---- gaussian blur (p=0.2 * p_per_channel 0.5, sigma U(0.5,1.0)/axis) ---- + if rng.random_sample() < 0.1: + sigma = (float(rng.uniform(0.5, 1.0)), float(rng.uniform(0.5, 1.0)), + float(rng.uniform(0.5, 1.0))) + x = ndimage.gaussian_filter(x, sigma=sigma) + + # ---- multiplicative brightness (p=0.15, U(0.75,1.25)) ---- + if rng.random_sample() < 0.15: + x = x * float(rng.uniform(0.75, 1.25)) + + # ---- contrast (p=0.15, factor U(0.75,1.25) straddled, preserve range) ---- + if rng.random_sample() < 0.15: + c = _bg_contrast_sample(rng, 0.75, 1.25) + mn = float(x.mean()) + minm, maxm = float(x.min()), float(x.max()) + x = (x - mn) * c + mn + x = np.clip(x, minm, maxm) + + # ---- simulate low resolution (p=0.25, zoom U(0.5,1), down order1 up order0) ---- + if rng.random_sample() < 0.25: + s = float(rng.uniform(0.5, 1.0)) + xz = ndimage.zoom(ndimage.zoom(x, s, order=1, mode='constant', cval=0.0), + 1.0 / s, order=0, mode='constant', cval=0.0) + x = _center_crop_or_pad(xz, x.shape) # double-zoom truncation shifts size + + # ---- gamma (p=0.1 inverted + p=0.3 plain; retain stats) ---- + for p, invert in ((0.1, True), (0.3, False)): + if rng.random_sample() < p: + g = _bg_contrast_sample(rng, 0.7, 1.5) + if invert: + x = -x + mn = float(x.mean()) + sd = float(x.std()) + minm = float(x.min()) + rnge = float(x.max() - minm) + x = np.power((x - minm) / (rnge + 1e-7), g) * rnge + minm + x = (x - x.mean()) / (x.std() + 1e-8) * sd + mn + if invert: + x = -x + + # ---- mirror (3 axes, p=0.5) ---- + for ax in range(3): + if rng.random_sample() < 0.5: + x = np.flip(x, ax).copy() + if lab is not None: + lab = np.flip(lab, ax).copy() + + out = x[np.newaxis].astype(np.float32, copy=False) + return out, (None if lab is None else lab.astype(np.int64)) + + +def zscore_volume(vol): + """Per-volume Z-score over ALL voxels (nnU-Net ZScoreNormalization, no mask).""" + v = np.asarray(vol, dtype=np.float32) + mu = v.mean() + sd = v.std() + if sd < 1e-8: + sd = 1.0 + return ((v - mu) / sd).astype(np.float32) + + +def downsample_label_nn(lab, factor=2): + """Nearest-neighbor label downsampling by `factor` per axis (B's + DownsampleSegForDSTransform, order=0).""" + out = ndimage.zoom(lab.astype(np.float32), 1.0 / factor, order=0, + mode='constant', cval=0) + return np.rint(out).astype(np.int64) \ No newline at end of file diff --git a/src/dataset.py b/src/dataset.py index fe1ed38..b6bb640 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -29,14 +29,24 @@ class PatchDataset(Dataset): Pipeline C's protocol. Uniform-only random crops put a ~20 mm tumor blob at a random edge position in most 96^3 patches and starve the loss of foreground signal on sparse targets. Unlabeled rows use plain uniform crops. + + b_aligned options (Pipeline B's data protocol, see src/augment_b.py): + zscore=True per-volume Z-score over all voxels (B's ZScoreNormalization, + no mask); the [0, 1.5] percentile clip is then skipped + b_aug=True B's full augmentation (rot/zoom/noise/blur/brightness/ + contrast/lowres/gamma/mirror) instead of the legacy set """ - def __init__(self, rows, patch=(96, 96, 96), train=True, seed=0, pos_neg=0.5): + def __init__(self, rows, patch=(96, 96, 96), train=True, seed=0, pos_neg=0.5, + zscore=False, b_aug=False, no_aug=False): self.rows = list(rows) self.patch = tuple(patch) self.train = train self.pos_neg = pos_neg self.seed = seed + self.zscore = zscore + self.b_aug = b_aug + self.no_aug = no_aug # crop but do not augment (B-style val patches) self._rng_wid = None self._rng = None @@ -48,7 +58,11 @@ class PatchDataset(Dataset): w = float(row.get("w", 1.0)) img = _read_nii(row["pimg"]).astype(np.float32) img = np.nan_to_num(img, nan=0.0, posinf=1.0, neginf=0.0) - img = np.clip(img, 0.0, 1.5) + if self.zscore: + from src.augment_b import zscore_volume + img = zscore_volume(img) # B's per-volume Z-score; no [0,1.5] clip + else: + img = np.clip(img, 0.0, 1.5) p = tuple(max((img.shape[i] + 7) // 8 * 8, self.patch[i] if self.train else 8) for i in range(3)) if any(a < b for a, b in zip(img.shape, p)): img = np.pad(img, [(0, b - a) for a, b in zip(img.shape, p)]) @@ -77,7 +91,14 @@ class PatchDataset(Dataset): img = img[s0[0]:s0[0] + self.patch[0], s0[1]:s0[1] + self.patch[1], s0[2]:s0[2] + self.patch[2]] if label is not None: label = label[s0[0]:s0[0] + self.patch[0], s0[1]:s0[1] + self.patch[1], s0[2]:s0[2] + self.patch[2]] - img, label = self._augment(img, label) + if self.no_aug: + pass # val-style: crop only + elif self.b_aug: + from src.augment_b import augment_b + img, label = augment_b(img[np.newaxis], label, self._worker_rng()) + img = img[0] + else: + img, label = self._augment(img, label) else: if any(img.shape[i] < self.patch[i] for i in range(3)): img = np.pad(img, [(0, max(self.patch[i] - img.shape[i], 0)) for i in range(3)]) @@ -159,8 +180,9 @@ def _collate(items): def make_dataloader(rows, patch, batch_size, train, num_workers=8, seed=0, persistent=True, - pos_neg=0.5): - ds = PatchDataset(rows, patch=patch, train=train, seed=seed, pos_neg=pos_neg) + pos_neg=0.5, zscore=False, b_aug=False, no_aug=False): + ds = PatchDataset(rows, patch=patch, train=train, seed=seed, pos_neg=pos_neg, + zscore=zscore, b_aug=b_aug, no_aug=no_aug) dl = DataLoader(ds, batch_size=batch_size, shuffle=train, num_workers=num_workers, collate_fn=_collate, drop_last=train and len(ds) > batch_size, pin_memory=True, persistent_workers=(train and num_workers > 0 and persistent), diff --git a/src/losses.py b/src/losses.py index 71463c8..af6dece 100644 --- a/src/losses.py +++ b/src/losses.py @@ -53,4 +53,71 @@ def per_sample_loss(logits, target, dice_weight=0.5, fg_only=False, fg_ce_weight def one_hot(label, num_classes=2): one = torch.zeros(label.size(0), num_classes, *label.shape[1:], device=label.device, dtype=torch.long) - return one.scatter_(1, label.unsqueeze(1), 1) \ No newline at end of file + return one.scatter_(1, label.unsqueeze(1), 1) + + +def nn_dc_and_ce(logits, cls): + """nnU-Net-style DC_and_CE for the 2-class problem (see the B pipeline). + + Replicates nnunetv2's DC_and_CE_loss + SoftDiceLoss as configured by the + 3d_fullres plans for this dataset: + - CE over class indices, unweighted (RobustCrossEntropyLoss) + - soft dice with do_bg=True (background and tumor, averaged) and + batch_dice=True (tp/fp/fn aggregated over the WHOLE batch, not + per-sample means) + - smooth = 1.0, returned as -dice, so loss = CE - mean_class(dice) + """ + ce = F.cross_entropy(logits, cls) + p = F.softmax(logits, dim=1) + g = F.one_hot(cls, p.size(1)).permute(0, -1, *range(1, cls.ndim)).float() + axes = (0, *range(2, p.dim())) + tp = (p * g).sum(axes) + fp = (p * (1.0 - g)).sum(axes) + fn = ((1.0 - p) * g).sum(axes) + dc = (2.0 * tp + 1.0) / (2.0 * tp + fp + fn + 1.0) + return ce - dc.mean() + + +def b_dc_and_ce(logits, cls, do_bg=False, smooth=1e-5): + """Pipeline B's ACTUAL loss, read from the installed nnunetv2 (2.8.1): + DC_and_CE_loss(soft_dice_kwargs={'batch_dice': True, 'smooth': 1e-5, + 'do_bg': False}, ce_kwargs={}, weight_ce=1, weight_dice=1) with + MemoryEfficientSoftDiceLoss. + + i.e. unweighted CE + batch-level soft dice on the TUMOR class only + (do_bg=False; the background class is NOT averaged in), smooth=1e-5, + returned as CE - dice. This differs from nn_dc_and_ce above, which was + an earlier (do_bg=True, smooth=1.0) approximation of B's loss. + """ + ce = F.cross_entropy(logits, cls) + p = F.softmax(logits, dim=1) # (B, C, z, y, x) + one = F.one_hot(cls, p.size(1)).permute(0, -1, *range(1, cls.ndim)).float() + if not do_bg: + p = p[:, 1:] + one = one[:, 1:] + axes = tuple(range(2, p.ndim)) # spatial -> (B, C) + inter = (p * one).sum(axes) + sum_pred = p.sum(axes) + sum_gt = one.sum(axes) + inter, sum_pred, sum_gt = inter.sum(0), sum_pred.sum(0), sum_gt.sum(0) # batch dice + dc = (2.0 * inter + smooth) / (sum_gt + sum_pred + smooth).clamp_min(1e-8) + return ce - dc.mean() + + +def deep_sup_loss(outputs, cls_list, loss_f, weights=(4.0/7.0, 2.0/7.0, 1.0/7.0)): + """nnU-Net DeepSupervisionWrapper: weight-scaled sum of per-scale losses. + + outputs: list of logit tensors, highest resolution first. + cls_list: matching downsampled class targets (same order). + B's convention (3d_fullres): raw weights [1, 1/2, 1/4, 1/8] with the lowest + scale (bottleneck) set to 0 and the rest renormalized to sum 1 -> + [4/7, 2/7, 1/7, 0]. A/C expose the 3 used scales (1x, 1/2, 1/4), so the + default here is (4/7, 2/7, 1/7). + """ + assert len(outputs) == len(cls_list) == len(weights) + loss = outputs[0].new_zeros(()) + for o, c, w in zip(outputs, cls_list, weights): + if w == 0: + continue + loss = loss + w * loss_f(o, c) + return loss \ No newline at end of file diff --git a/src/training.py b/src/training.py index aa1ec66..e0b0d9d 100644 --- a/src/training.py +++ b/src/training.py @@ -9,14 +9,18 @@ from src.dataset import make_dataloader from src.losses import per_sample_loss, one_hot -def build_model(base=16, out_ch=2, device="cuda"): - return Unet3D(in_ch=1, out_ch=out_ch, base=base, depth=4).to(device) +def build_model(base=16, out_ch=2, device="cuda", leaky=False): + return Unet3D(in_ch=1, out_ch=out_ch, base=base, depth=4, leaky=leaky).to(device) @torch.no_grad() -def sliding_window_probs(model, vol, device, win=96, step=48, tta=True): - """vol: (H,W,D) float32 -> tumor probability (H,W,D) numpy float32.""" +def sliding_window_probs(model, vol, device, win=96, step=48, tta=True, zscore=False): + """vol: (H,W,D) float32 -> tumor probability (H,W,D) numpy float32. + zscore=True matches a B-aligned model's input (per-volume Z-score).""" model.eval() + if zscore: + from src.augment_b import zscore_volume + vol = zscore_volume(vol) h, w, d = vol.shape vol = np.pad(vol, [(0, max(h, win) - h), (0, max(w, win) - w), (0, max(d, win) - d)]) h, w, d = vol.shape @@ -43,8 +47,13 @@ def sliding_window_probs(model, vol, device, win=96, step=48, tta=True): for cy in cs[1]: for cx in cs[0]: sl = (slice(cx, cx + win), slice(cy, cy + win), slice(cz, cz + win)) + patch0 = vol[sl] for tf in flips: - patch = tf(vol)[sl] + # flip the PATCH, not the whole volume: tf(vol)[sl] would window the + # flipped volume at the original grid position (the mirrored window), + # so the views' probabilities would be averaged at the wrong location + # and collapse to ~1/n_flips at off-mirror targets (TTA blending bug) + patch = tf(patch0) t = torch.from_numpy(patch).float().unsqueeze(0).unsqueeze(0).to(device) with torch.autocast("cuda", dtype=torch.bfloat16): logits = model(t) diff --git a/src/unet3d.py b/src/unet3d.py index 3fa6243..8353db5 100644 --- a/src/unet3d.py +++ b/src/unet3d.py @@ -3,35 +3,64 @@ import torch.nn as nn import torch.nn.functional as F +def make_norm(ch, norm="batch"): + if norm == "instance": + # Pipeline B's norm: InstanceNorm3d(eps=1e-5, affine=True) + return nn.InstanceNorm3d(ch, affine=True, eps=1e-5) + return nn.BatchNorm3d(ch) + + class DoubleConv(nn.Module): - def __init__(self, in_ch, out_ch): + def __init__(self, in_ch, out_ch, leaky=False, norm="batch"): super().__init__() + act = nn.LeakyReLU(0.01, inplace=True) if leaky else nn.ReLU(inplace=True) self.block = nn.Sequential( nn.Conv3d(in_ch, out_ch, 3, padding=1), - nn.BatchNorm3d(out_ch), - nn.ReLU(inplace=True), + make_norm(out_ch, norm), + act, nn.Conv3d(out_ch, out_ch, 3, padding=1), - nn.BatchNorm3d(out_ch), - nn.ReLU(inplace=True), + make_norm(out_ch, norm), + act, ) def forward(self, x): return self.block(x) +def init_b(module): + """Pipeline B's network init (nnunetv2 InitWeights_He(1e-2)): + kaiming_normal_(a=0.01) on every conv weight, zero bias.""" + if isinstance(module, (nn.Conv3d, nn.ConvTranspose3d, nn.Conv2d)): + nn.init.kaiming_normal_(module.weight, a=0.01) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + class Unet3D(nn.Module): - def __init__(self, in_ch=1, out_ch=2, base=16, depth=4): + def __init__(self, in_ch=1, out_ch=2, base=16, depth=4, leaky=False, norm="batch", + deep_sup=False, b_init=False): super().__init__() chs = [base * 2 ** i for i in range(depth + 1)] - self.encs = nn.ModuleList([DoubleConv(in_ch if i == 0 else chs[i - 1], chs[i]) for i in range(depth)]) + self.encs = nn.ModuleList([DoubleConv(in_ch if i == 0 else chs[i - 1], chs[i], leaky, norm) + for i in range(depth)]) self.pools = nn.ModuleList([nn.MaxPool3d(2) for _ in range(depth - 1)]) self.decs = nn.ModuleList([ - DoubleConv(chs[depth - 1 - k] + chs[depth - 2 - k], chs[depth - 2 - k]) + DoubleConv(chs[depth - 1 - k] + chs[depth - 2 - k], chs[depth - 2 - k], leaky, norm) for k in range(depth - 1) ]) self.head = nn.Conv3d(chs[0], out_ch, 1) self.depth = depth self.chs = chs + self.deep_sup = deep_sup + if deep_sup: + # B's deep supervision: extra 1x1 heads at the 1/2 and 1/4 decoder + # scales (loss weights 4/7, 2/7, 1/7 for 1x/1/2/1/4; the 1/8 + # bottleneck is dropped, as nnU-Net sets the lowest scale to 0). + # dec k=1 outputs chs[1] @ 1/2 res, k=0 outputs chs[2] @ 1/4 res. + self.aux_half = nn.Conv3d(chs[1], out_ch, 1) + self.aux_quarter = nn.Conv3d(chs[2], out_ch, 1) + if b_init: + self.apply(init_b) def forward(self, x): skips = [] @@ -41,8 +70,19 @@ class Unet3D(nn.Module): skips.append(h) if i < len(self.pools): h = self.pools[i](h) + aux_half = aux_quarter = None for k, dec in enumerate(self.decs): h = F.interpolate(h, scale_factor=2, mode="nearest") h = torch.cat([h, skips[self.depth - 2 - k]], dim=1) h = dec(h) - return self.head(h) \ No newline at end of file + if self.deep_sup: + if k == 1: + aux_half = h + elif k == 0: + aux_quarter = h + out = self.head(h) + if self.deep_sup and self.training: + # B's deep supervision: (final, aux@1/2, aux@1/4); in eval mode the + # aux heads are unused so a plain tensor is returned (inference-safe) + return out, self.aux_half(aux_half), self.aux_quarter(aux_quarter) + return out \ No newline at end of file