feat(benchmark): update benchmarking scripts with enhanced loss functions and dataset handling
This commit is contained in:
parent
2453fb48ba
commit
99bea82841
5 changed files with 210 additions and 44 deletions
103
BENCHMARKING.md
103
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 <total>`).
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue