diff --git a/BENCHMARKING.md b/BENCHMARKING.md index 1e75559..8789517 100644 --- a/BENCHMARKING.md +++ b/BENCHMARKING.md @@ -83,35 +83,98 @@ ntuh). Artifacts: `results/benchmark_pipelines_smoke.json`, ModuleNotFoundError) — installed `pip install monai==1.6.1 --no-deps` (torch 2.14+cu126 untouched); `import monai` verified. -## Full-dataset 10-epoch run (2026-10-04) — interim +## Full-dataset 10-epoch run (2026-10-05, updated dataset) -`--train-epochs 10`, default budgets (2856 train rows = 80% of split_train, -42 eval rows = 10% of split_test, TTA on), A/C/GPU0·GPU2 finished, B still -preprocessing at time of writing: +`--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/`). ``` -python scripts/benchmark_pipelines.py --pipelines A --gpu 0 --train-epochs 10 \ - --out results/benchmark_full_A.json # done 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 # in progress (nnU-Net prep) + --out results/benchmark_full_B.json # done (local-disk NNU_ROOT, see below) ``` | pipeline | train s/epoch | wall | tr VRAM | s/vol (TTA) | in VRAM | params | weights | dice @10ep | |---|---|---|---|---|---|---|---|---| -| A Unet3D(16, depth 4) | 258.8 (952 steps) | 46 min | 2.0 GB | 3.51 ± 0.93 | 0.5 GB | 1.46 M | 17.6 MB | **0.0008** | -| C MONAI UNet(16→128) | 284.3 (952 steps) | 48 min | 0.4 GB | 0.38 ± 0.06 | 0.5 GB | 1.19 M | 14.3 MB | **0.2119** | -| B nnU-Net 3d_fullres | (pending) | (pending) | — | — | — | 30.8 M | 246 MB | (pending) | +| 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) | -Dice distribution over 41 timed test vols: A p25/p75 = 0.0000/0.0000 (98% ≈ 0); -C p25/p75 = 0.0000/0.4048 (29% ≈ 0). Training loss (epoch 1 → 10): A -0.335 → 0.209, C 0.403 → 0.188. 10 epochs is far below convergence for these -sparse targets (median tumor 3.1 k mm³); the A/C numbers here are -intermediate, and C is already clearly ahead at this budget. Checkpoints at +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 `). -Note: B's first `plan_and_preprocess` attempt died and the built-in 60 s retry -re-ran it successfully. + +### A: why dice ≈ 0, and the fix (2026-10-05) + +The A baseline predicts all-background (38/39 test rows dice exactly 0.0). +Root-cause work: + +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. + +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`. + +### B: CIFS failures and local-disk NNU_ROOT (2026-10-05) + +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: + +- `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. + +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). ## Resuming training (`--resume`) @@ -144,9 +207,9 @@ python scripts/benchmark_pipelines.py --pipelines "A B C" --resume \ ## Readiness notes for a full run - Budget knobs: `--train-epochs` (default 100), `--max-rows` (default 80% of - split_train ≈ 2753 for A/C and for B's `--nnu-cases`, post-exclusion; - ≈ 2856 pre-exclusion), `--max-eval-rows` (default 10% of split_test ≈ 40 - post-exclusion), `--no-tta`. + 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 diff --git a/scripts/benchmark_pipelines.py b/scripts/benchmark_pipelines.py index 6ed18fe..f326893 100644 --- a/scripts/benchmark_pipelines.py +++ b/scripts/benchmark_pipelines.py @@ -46,6 +46,7 @@ import subprocess import sys import threading import time +import traceback from datetime import datetime ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -80,6 +81,11 @@ def parse_args(): ap.add_argument("--workers", type=int, default=4, help="dataloader workers (A/C)") ap.add_argument("--batch", type=int, default=3, help="per-GPU batch (A/C)") ap.add_argument("--patch", type=int, default=96, help="patch/window (A/C)") + ap.add_argument("--a-fg-dice", action="store_true", + help="A: use tumor-class-only soft dice in the per-sample loss " + "(instead of the background+tumor mean)") + ap.add_argument("--a-fg-ce-weight", type=float, default=1.0, + help="A: tumor class weight in the CE part of the loss (1.0 = unweighted)") ap.add_argument("--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) " @@ -256,7 +262,8 @@ def bench_A(args, train_rows, test_rows): 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)) + 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) @@ -290,7 +297,7 @@ def bench_A(args, train_rows, test_rows): args.patch, step, tta) torch.cuda.synchronize(dev) reset_peak(dev) - times, dices = [], [] + times, dices, pmaxs = [], [], [] with torch.no_grad(): for r in test_rows[args.eval_warmup:]: vol = read_nii_arr(r["pimg"]).astype("float32") @@ -302,6 +309,7 @@ def bench_A(args, train_rows, test_rows): # 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] @@ -311,8 +319,9 @@ def bench_A(args, train_rows, test_rows): res["infer_vram_mb"] = peak_mb(dev) res["dice"] = round(float(np.mean(dices)), 4) if dices else 0.0 res["dice_per_row"] = [{"key": test_rows[args.eval_warmup + i]["key"], - "dice": round(dices[i], 4), "s": round(times[i], 2)} - for i in range(len(dices))] + "dice": round(dices[i], 4), "s": round(times[i], 2), + "pmax": round(float(pm.max()), 4)} + for i, pm in zip(range(len(dices)), pmaxs)] del model free_gpu() return res @@ -463,7 +472,8 @@ def bench_B(args, train_rows, test_rows): res["raw_build_s"] = round(time.time() - t0, 1) if not args.nnu_skip_prep: t0 = time.time() - N.run(N.plan_preprocess_cmd(8), os.path.join(d("logs/nnu"), "bench_plan.log"), env=N.nnu_env()) + N.run(N.plan_preprocess_cmd(8), os.path.join(d("logs/nnu"), "bench_plan.log"), env=N.nnu_env(), + retries=5) # CIFS /mnt/b4 drops during the ~3h preprocess; extra attempts are cheap res["plan_preprocess_s"] = round(time.time() - t0, 1) else: res["plan_preprocess_s"] = 0.0 @@ -518,8 +528,9 @@ def bench_B(args, train_rows, test_rows): te_keys = [r["key"] for r in te] te_native = native_rows_for(te_keys) - in_dir = d("runs_bench/B/pred_in") - out_dir = d("runs_bench/B/pred_out") + pred_base = os.environ.get("BENCH_B_DIR", os.path.join(d("runs_bench"), "B")) + in_dir = os.path.join(pred_base, "pred_in") + out_dir = os.path.join(pred_base, "pred_out") shutil.rmtree(in_dir, ignore_errors=True) shutil.rmtree(out_dir, ignore_errors=True) os.makedirs(in_dir) @@ -669,8 +680,9 @@ def main(): r = runners[p](args, train_rows, test_rows) r["wall_s"] = round(time.time() - t0, 1) except Exception as e: # noqa - r = {"pipeline": p, "error": repr(e), "wall_s": round(time.time() - t0, 1)} - print(f"[bench] pipeline {p} FAILED: {e!r}") + tb = traceback.format_exc() + r = {"pipeline": p, "error": repr(e), "traceback": tb, "wall_s": round(time.time() - t0, 1)} + print(f"[bench] pipeline {p} FAILED:\n{tb}") results.append(r) free_gpu() diff --git a/scripts_nnu/nnu_common.py b/scripts_nnu/nnu_common.py index 17faa23..4dc139e 100644 --- a/scripts_nnu/nnu_common.py +++ b/scripts_nnu/nnu_common.py @@ -31,7 +31,11 @@ TRAINERS_DIR = os.path.join(SCRIPTS_NNU, "trainers") def nnu_root(): - return d(os.path.join("nnu")) + # NNU_ROOT override: the CIFS /mnt/b4 mount drops writes transiently + # ("Host is down", "Error while setting the slice"), which kills the ~3h + # plan&preprocess and the multi-day train. Point it at local disk + # (e.g. /home/xfr/longitudinal_nnu) to keep all nnU-Net I/O off CIFS. + return os.environ.get("NNU_ROOT", d(os.path.join("nnu"))) def raw_ds(): @@ -156,14 +160,29 @@ def native_label_target(img_path, lab_path, link_path): sitk.WriteImage(out, link_path, True) +def rmtree_verified(path, tries=10, pause=2.0): + """Remove a tree and verify it is actually gone. CIFS (/mnt/b4, soft mount) + can transiently fail individual unlinks; a silent ignore_errors rmtree leaves + stale entries behind and the next rebuild's os.symlink dies with EEXIST.""" + for i in range(tries): + if not os.path.isdir(path): + return + shutil.rmtree(path, ignore_errors=True) + if not os.path.isdir(path) or len(os.listdir(path)) == 0: + return + time.sleep(pause * (i + 1)) + raise RuntimeError(f"could not fully remove {path} " + f"({len(os.listdir(path)) if os.path.isdir(path) else '?'} entries left)") + + def make_raw_dataset(rows, channel="0000"): """rows: list of {key, img, label} — img/label are native (unprocessed) NIfTI paths. Rebuilds imagesTr/labelsTr + dataset.json. Images are symlinked as-is; labels via native_label_target. See section note.""" img_dir = os.path.join(raw_ds(), "imagesTr") lab_dir = os.path.join(raw_ds(), "labelsTr") - shutil.rmtree(img_dir, ignore_errors=True) - shutil.rmtree(lab_dir, ignore_errors=True) + rmtree_verified(img_dir) + rmtree_verified(lab_dir) os.makedirs(img_dir, exist_ok=True) os.makedirs(lab_dir, exist_ok=True) seen = set() @@ -176,8 +195,13 @@ def make_raw_dataset(rows, channel="0000"): raise FileNotFoundError(f"image missing for {key}: {img}") if not os.path.exists(lab): raise FileNotFoundError(f"label missing for {key}: {lab}") - os.symlink(os.path.abspath(img), os.path.join(img_dir, f"{key}_{channel}.nii.gz")) - native_label_target(img, lab, os.path.join(lab_dir, f"{key}.nii.gz")) + img_link = os.path.join(img_dir, f"{key}_{channel}.nii.gz") + lab_link = os.path.join(lab_dir, f"{key}.nii.gz") + for lp in (img_link, lab_link): # belt-and-braces against stale CIFS state + if os.path.lexists(lp): + os.unlink(lp) + os.symlink(os.path.abspath(img), img_link) + native_label_target(img, lab, lab_link) ds_json = {"channel_names": {"0": "t1c"}, "labels": {"background": 0, "tumor": 1}, "numTraining": len(seen), diff --git a/src/dataset.py b/src/dataset.py index b205af0..fe1ed38 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -21,13 +21,24 @@ class PatchDataset(Dataset): row: {key, pimg, plabel?, is_neg?, w (loss weight, default 1.0)} yields (img(1,H,W,D), label(H,W,D) or None, weight) + + Training crops for labeled rows: `pos_neg` fraction are foreground-centered + (random tumor voxel as crop center), the rest are foreground-free background + crops (up to 20 random retries, plain uniform if the volume is too full to + find one) — mirrors MONAI RandCropByPosNegLabeld(pos=1, neg=1), which is + 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. """ - def __init__(self, rows, patch=(96, 96, 96), train=True, seed=0): + def __init__(self, rows, patch=(96, 96, 96), train=True, seed=0, pos_neg=0.5): self.rows = list(rows) self.patch = tuple(patch) self.train = train - self.rng = np.random.RandomState(seed) + self.pos_neg = pos_neg + self.seed = seed + self._rng_wid = None + self._rng = None def __len__(self): return len(self.rows) @@ -62,8 +73,7 @@ class PatchDataset(Dataset): raise RuntimeError(f"label file unreadable for {row.get('key')}: {row['plabel']}") if self.train: - s0 = [self.rng.randint(0, img.shape[i] - self.patch[i] + 1) if img.shape[i] >= self.patch[i] else 0 - for i in range(3)] + s0 = self._crop_origin(img.shape, label) 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]] @@ -78,6 +88,41 @@ class PatchDataset(Dataset): lab = torch.from_numpy(label) if label is not None else None return t, lab, torch.tensor(w, dtype=torch.float32) + def _worker_rng(self): + # forked dataloader workers inherit the parent's rng state; reseed per + # worker so each one draws an independent crop-position stream + info = torch.utils.data.get_worker_info() + wid = info.id if info is not None else 0 + if self._rng_wid != wid: + self._rng = np.random.RandomState(self.seed + 104729 * wid) + self._rng_wid = wid + return self._rng + + def _crop_origin(self, shape, label): + """Crop origin: fg-centered for `pos_neg` of labeled rows, otherwise a + foreground-free background crop (20 retries, then plain uniform).""" + rng = self._worker_rng() + win = self.patch + + def uniform(): + return [int(rng.randint(0, n - w + 1)) if n >= w else 0 for n, w in zip(shape, win)] + + has_fg = label is not None and int(label.sum()) > 0 + if has_fg and rng.rand() < self.pos_neg: + idx = np.argwhere(label > 0) + c = idx[int(rng.randint(len(idx)))] + s = [int(c[i] - win[i] // 2) for i in range(3)] + return [min(max(s[i], 0), shape[i] - win[i]) for i in range(3)] + s = uniform() + if has_fg: + sl = tuple(slice(s[i], s[i] + win[i]) for i in range(3)) + for _ in range(20): + if int(label[sl].sum()) == 0: + return s + s = uniform() + sl = tuple(slice(s[i], s[i] + win[i]) for i in range(3)) + return s + @staticmethod def _augment(img, label): if np.random.rand() > 0.5: @@ -113,8 +158,9 @@ def _collate(items): return imgs, labs, wts -def make_dataloader(rows, patch, batch_size, train, num_workers=8, seed=0, persistent=True): - ds = PatchDataset(rows, patch=patch, train=train, seed=seed) +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) dl = DataLoader(ds, batch_size=batch_size, shuffle=train, num_workers=num_workers, collate_fn=_collate, drop_last=train and len(ds) > batch_size, pin_memory=True, persistent_workers=(train and num_workers > 0 and persistent), diff --git a/src/losses.py b/src/losses.py index 8ea156c..71463c8 100644 --- a/src/losses.py +++ b/src/losses.py @@ -22,11 +22,32 @@ class SoftDiceLoss(nn.Module): return torch.stack(losses, dim=0).mean(dim=0) -def per_sample_loss(logits, target, dice_weight=0.5): - """target: (B,C,...) one-hot -> per-sample loss (B,)""" - bce = F.cross_entropy(logits, target.argmax(dim=1), reduction="none") +def soft_dice_fg(logits, target, eps=1e-5): + """Soft dice loss on the tumor class (channel 1) only. Averaging the dice + over background+tumor dilutes the foreground signal ~10^4-fold on sparse + targets; the tumor-only variant keeps the same 0.5/0.5 CE+Dice blend but + points all of the Dice term at the tumor.""" + p = F.softmax(logits, dim=1)[:, 1] + t = target[:, 1].float() + dim = tuple(range(1, p.dim())) + inter = (p * t).sum(dim) + denom = p.sum(dim) + t.sum(dim) + return 1 - (2 * inter + eps) / (denom + eps) + + +def per_sample_loss(logits, target, dice_weight=0.5, fg_only=False, fg_ce_weight=1.0): + """target: (B,C,...) one-hot -> per-sample loss (B,) + + fg_ce_weight: class weight [1, w] on the CE part — multiplies the CE on + tumor voxels. Uniform CE on a ~0.05% target leaves the model + under-confident (p_tumor maxes at ~0.4 at a converged loss); weighting the + tumor class sharpens saturation at the cost of more false-positive pressure, + countered by the dice term.""" + cls = target.argmax(dim=1) + w = None if fg_ce_weight == 1.0 else torch.tensor([1.0, float(fg_ce_weight)], device=logits.device) + bce = F.cross_entropy(logits, cls, weight=w, reduction="none") bce = bce.mean(dim=tuple(range(1, bce.dim()))) - dcl = SoftDiceLoss()(logits, target) + dcl = soft_dice_fg(logits, target) if fg_only else SoftDiceLoss()(logits, target) return (1 - dice_weight) * bce + dice_weight * dcl