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.
This commit is contained in:
parent
5aabe950bd
commit
e7721760f2
20 changed files with 1584 additions and 446 deletions
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -11,3 +11,5 @@ __pycache__/
|
|||
elastix/
|
||||
niftyreg/NiftyReg-Ubuntu-v2.1.1.zip
|
||||
niftyreg/NiftyReg-Ubuntu-v2.1.1/
|
||||
|
||||
.kilo/
|
||||
500
BENCHMARKING.md
500
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/<key>{,_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/<key>{,_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 <total>`).
|
||||
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).
|
||||
- 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.
|
||||
51
README.md
51
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 \
|
||||
|
|
|
|||
|
|
@ -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/<pipeline>__<case>.png` (9 per
|
||||
`results/bench_reg/bench_registration_affine/qa/<pipeline>__<case>.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
|
||||
|
|
|
|||
|
|
@ -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/<pipeline>__<case>.png`
|
||||
`results/bench_reg/bench_registration_deformable/qa/<pipeline>__<case>.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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 <out>.partial.json (crash protection:
|
||||
mid-training evals + per-epoch loss survive a failed run; atomic replace)."""
|
||||
try:
|
||||
out_f = args.out or os.path.join(ROOT, "results", "benchmark_pipelines", "benchmark_pipelines.json")
|
||||
payload = {"timestamp": datetime.now().isoformat(timespec="seconds"), "partial": True,
|
||||
"settings": {k: v for k, v in vars(args).items()},
|
||||
"results": [res]}
|
||||
tmp = out_f + ".partial.json.tmp"
|
||||
with open(tmp, "w") as f:
|
||||
json.dump(payload, f, indent=1)
|
||||
os.replace(tmp, out_f + ".partial.json")
|
||||
except Exception as e: # noqa
|
||||
print(f"partial dump failed: {e!r}", flush=True)
|
||||
|
||||
|
||||
def _reconstructed_lr(steps_per_epoch, epochs_total, start_epoch, base_lr=3e-4, eta_frac=0.05):
|
||||
"""LR at the start of each epoch for the A/C schedule (3e-4, 300-step
|
||||
LinearLR warmup from 0.1x, then CosineAnnealing to base_lr*eta_frac over
|
||||
the TOTAL run's steps) — used to draw the lr panel for runs that predate
|
||||
per-epoch LR logging."""
|
||||
total = steps_per_epoch * epochs_total
|
||||
warmup = min(300, max(10, total // 10))
|
||||
eta = base_lr * eta_frac
|
||||
out = []
|
||||
for ep in range(start_epoch + 1, epochs_total + 1):
|
||||
s = (ep - 1) * steps_per_epoch
|
||||
if s < warmup:
|
||||
lr = base_lr * (0.1 + 0.9 * s / max(warmup, 1))
|
||||
else:
|
||||
t = min(s - warmup, max(total - warmup, 1))
|
||||
lr = eta + 0.5 * (base_lr - eta) * (1 + math.cos(math.pi * t / max(total - warmup, 1)))
|
||||
out.append(lr)
|
||||
return out
|
||||
|
||||
|
||||
def plot_progress(r, png):
|
||||
"""nnU-Net-style progress figure (like fold_0/progress.png): stacked panels
|
||||
1) train loss (left) + test dice (right axis, from res['evals']),
|
||||
2) epoch duration (when per-epoch times are logged),
|
||||
3) learning rate (logged, or reconstructed for the A/C schedule)."""
|
||||
try:
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
try:
|
||||
plt.style.use("seaborn-v0_8-darkgrid")
|
||||
except OSError:
|
||||
plt.style.use("darkgrid")
|
||||
except Exception as e: # noqa
|
||||
print("progress plot failed:", repr(e))
|
||||
return
|
||||
loss = r.get("per_epoch_loss") or []
|
||||
start = int(r.get("resumed_from_epoch", 0)) + 1
|
||||
epochs = list(range(start, start + len(loss)))
|
||||
ev = r.get("evals") or {}
|
||||
mid = sorted((int(t.split("_")[1]), d["dice"]) for t, d in ev.items() if t.startswith("epoch_"))
|
||||
final_ep = int(r.get("epochs_total", epochs[-1] if epochs else (mid[-1][0] if mid else 0)))
|
||||
dice_pts = mid + ([(final_ep, ev["final"]["dice"])] if "final" in ev else [])
|
||||
dur = r.get("per_epoch_time") or []
|
||||
if len(dur) != len(epochs):
|
||||
dur = []
|
||||
lr = r.get("per_epoch_lr") or []
|
||||
if len(lr) != len(epochs):
|
||||
lr = _reconstructed_lr(int(r.get("steps_per_epoch", 0)), final_ep, start - 1) \
|
||||
if r.get("steps_per_epoch") and final_ep else []
|
||||
have = [p for p in (("loss", bool(loss)), ("dur", bool(dur)), ("lr", bool(lr))) if p[1]]
|
||||
if not have:
|
||||
return
|
||||
blue, green = "#1f77b4", "#2ca02c"
|
||||
fig, axes = plt.subplots(len(have), 1, figsize=(9, 4.8 * len(have)), sharex=True)
|
||||
if len(have) == 1:
|
||||
axes = [axes]
|
||||
for ax, (name, _) in zip(axes, have):
|
||||
if name == "loss":
|
||||
ax.plot(epochs, loss, color=blue, lw=1.0, label="loss_tr")
|
||||
ax.set_ylabel("loss")
|
||||
handles, labels = ax.get_legend_handles_labels()
|
||||
if dice_pts:
|
||||
ax2 = ax.twinx()
|
||||
ax2.plot([p[0] for p in dice_pts], [p[1] for p in dice_pts], color=green,
|
||||
lw=1.5, marker="o", ms=4, label="test dice (TTA)")
|
||||
ax2.set_ylabel("test dice")
|
||||
handles += ax2.get_legend_handles_labels()[0]
|
||||
labels += ax2.get_legend_handles_labels()[1]
|
||||
ax.legend(handles, labels, loc="upper left", fontsize=8)
|
||||
elif name == "dur":
|
||||
ax.plot(epochs, dur, color=blue, lw=1.0, label="epoch duration")
|
||||
ax.set_ylabel("time [s]")
|
||||
ax.legend(loc="upper left", fontsize=8)
|
||||
elif name == "lr":
|
||||
ax.plot(epochs, lr, color=blue, lw=1.2, label="learning rate")
|
||||
ax.set_ylabel("learning rate")
|
||||
ax.legend(loc="upper left", fontsize=8)
|
||||
ax.set_xlabel("epoch")
|
||||
fig.suptitle(f"pipeline {r['pipeline']} {r.get('backbone', '')} — {r.get('loss', '')}", fontsize=10)
|
||||
fig.tight_layout()
|
||||
fig.savefig(png, dpi=130)
|
||||
plt.close(fig)
|
||||
print("progress.png:", png)
|
||||
|
||||
|
||||
def plot(results, out_png, epochs):
|
||||
try:
|
||||
import matplotlib
|
||||
|
|
@ -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__":
|
||||
|
|
|
|||
37
scripts/run_b_100ep_stages.sh
Normal file
37
scripts/run_b_100ep_stages.sh
Normal file
|
|
@ -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) ==="
|
||||
19
scripts/run_c_bloss_100ep_after_b.sh
Normal file
19
scripts/run_c_bloss_100ep_after_b.sh
Normal file
|
|
@ -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)"
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
149
src/augment_b.py
Normal file
149
src/augment_b.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -54,3 +54,70 @@ 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)
|
||||
|
||||
|
||||
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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
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
|
||||
Loading…
Reference in a new issue