From a491ae05237ff3e2e28ac8f2576a36b7d81a9c35 Mon Sep 17 00:00:00 2001 From: Furen Xiao Date: Sat, 26 Sep 2026 03:44:12 +0800 Subject: [PATCH] feat(train): implement checkpointing and robust data loading Enhance the training pipeline with stateful checkpointing and improve the resilience of the data loading process against filesystem latency and transient I/O errors. - Implement auto-resuming in `train_ddp` by loading model, optimizer, and scheduler states from `state.pt`. - Add atomic state saving using temporary files to prevent corruption. - Introduce `_read_nii` with exponential backoff retries to handle transient NFS/filesystem failures during NIfTI reading. - Add explicit error handling for missing or unreadable label files in `PatchDataset`. - Update `sliding_window_probs` to conditionally apply Test-Time Augmentation (TTA) based on the `tta` parameter. - Add `scripts/test_dataloader.py` for verifying dataset integrity. --- scripts/test_dataloader.py | 24 +++++++++++++++++++ src/dataset.py | 48 ++++++++++++++++++++++++-------------- src/training.py | 33 +++++++++++++++++++++++--- 3 files changed, 84 insertions(+), 21 deletions(-) create mode 100644 scripts/test_dataloader.py diff --git a/scripts/test_dataloader.py b/scripts/test_dataloader.py new file mode 100644 index 0000000..9fec62b --- /dev/null +++ b/scripts/test_dataloader.py @@ -0,0 +1,24 @@ +"""Test: iterate the full train dataloader with workers, report None-label batches.""" +import os +import sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def main(): + from src import training + from src.common import load_jsonl + rows = load_jsonl(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "manifests", "split_train.jsonl")) + for seed in (0, 1, 2): + ds, dl = training.make_dataloader(rows, (96, 96, 96), 3, True, num_workers=4, seed=seed) + n_none, n = 0, 0 + for i, (img, lab, w) in enumerate(dl): + if lab is None: + n_none += 1 + print(" seed", seed, "batch", i, "None label", flush=True) + n += 1 + print("seed", seed, "batches", n, "none-label", n_none, flush=True) + print("TEST DONE") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/dataset.py b/src/dataset.py index 9883379..b205af0 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -1,13 +1,25 @@ import os +import time import numpy as np import torch from torch.utils.data import Dataset, DataLoader +def _read_nii(path, tries=10, delay=2.0): + import SimpleITK as sitk + for i in range(tries): + try: + return sitk.GetArrayFromImage(sitk.ReadImage(path)) + except Exception: + if i == tries - 1: + raise + time.sleep(delay * (i + 1)) + + class PatchDataset(Dataset): """Random 3D patches from preprocessed nifti volumes. - row: {key, pimg, plabel?, w (loss weight, default 1.0)} + 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) """ @@ -20,36 +32,34 @@ class PatchDataset(Dataset): def __len__(self): return len(self.rows) - @staticmethod - def _read(path): - import SimpleITK as sitk - return sitk.GetArrayFromImage(sitk.ReadImage(path)) - def __getitem__(self, idx): row = self.rows[idx] w = float(row.get("w", 1.0)) - img = self._read(row["pimg"]).astype(np.float32) + 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) - # pad to multiple of 8 and at least the patch size (training) / 256 max (eval irrelevant) 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)]) label = None - if row.get("plabel") and os.path.exists(row["plabel"]): - label = self._read(row["plabel"]).astype(np.int64) - # pad/trim label to the same shape as img + plabel = row.get("plabel") + if plabel: + if not os.path.exists(plabel): + for i in range(10): # transient NFS handle failures + time.sleep(2 * (i + 1)) + if os.path.exists(plabel): + break + else: + raise RuntimeError(f"label file vanished for {row.get('key')}: {plabel}") + label = _read_nii(plabel).astype(np.int64) if label.shape != img.shape: - lab_p = [] - for i in range(3): - pad_i = 0 - pad_j = max(img.shape[i] - label.shape[i], 0) - lab_p.append((pad_i, pad_j)) - label = np.pad(label, lab_p) + label = np.pad(label, [(0, max(img.shape[i] - label.shape[i], 0)) for i in range(3)]) label = label[: img.shape[0], : img.shape[1], : img.shape[2]] label = np.clip(np.rint(label).astype(np.int64), 0, 1) elif row.get("is_neg"): label = np.zeros(img.shape, dtype=np.int64) + elif row.get("plabel"): + 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 @@ -97,6 +107,7 @@ def _collate(items): if all(l is not None for l in labs): labs = torch.stack(labs) else: + # rows without labels (should not mix with labeled rows in a training batch) labs = None wts = torch.stack([it[2] for it in items]) return imgs, labs, wts @@ -106,5 +117,6 @@ def make_dataloader(rows, patch, batch_size, train, num_workers=8, seed=0, persi ds = PatchDataset(rows, patch=patch, train=train, seed=seed) 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), timeout=900) + persistent_workers=(train and num_workers > 0 and persistent), + timeout=0 if num_workers == 0 else 900) return ds, dl \ No newline at end of file diff --git a/src/training.py b/src/training.py index f5ed2c8..aa1ec66 100644 --- a/src/training.py +++ b/src/training.py @@ -31,9 +31,10 @@ def sliding_window_probs(model, vol, device, win=96, step=48, tta=True): cs = (coords(h), coords(w), coords(d)) acc = np.zeros((h, w, d), dtype=np.float32) cnt = np.zeros((h, w, d), dtype=np.float32) - flips = [lambda x: x, lambda x: x[:, ::-1].copy(), lambda x: x[:, :, ::-1].copy(), - lambda x: x[::-1, :, :].copy()] + flips = [lambda x: x] if tta: + flips += [lambda x: x[:, ::-1].copy(), lambda x: x[:, :, ::-1].copy(), lambda x: x[::-1, :, :].copy()] + if tta == 2: flips.append(lambda x: x[::-1, ::-1].copy()) flips.append(lambda x: x[::-1, :, ::-1].copy()) flips.append(lambda x: x[:, ::-1, ::-1].copy()) @@ -127,14 +128,37 @@ def train_ddp(rows, val_rows, args): ckpt_dir = args.ckpt_dir if rank == 0: os.makedirs(ckpt_dir, exist_ok=True) + start_epoch = 0 + state_f = os.path.join(ckpt_dir, "state.pt") + if world > 1: + dist.barrier() + if os.path.exists(state_f): + st = torch.load(state_f, map_location="cpu", weights_only=True) + model.load_state_dict(st["model"]) + opt.load_state_dict(st["opt"]) + sched.load_state_dict(st["sched"]) + best, best_epoch, start_epoch = st["best"], st["best_epoch"], st["epoch"] + if rank == 0: + print(f"[train] auto-resuming from state.pt @epoch {start_epoch} (best={best:.4f})", flush=True) + if rank == 0: print(f"[train] rows={len(rows)} val={len(val_rows)} epochs={args.epochs} steps/epoch={steps_per_epoch}", flush=True) + if world > 1: + dist.barrier() - for epoch in range(args.epochs): + def save_state(): + tmp = state_f + ".tmp" + torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "sched": sched.state_dict(), + "best": best, "best_epoch": best_epoch, "epoch": epoch + 1}, tmp) + os.replace(tmp, state_f) + + for epoch in range(start_epoch, args.epochs): if hasattr(dl.sampler, "set_epoch"): dl.sampler.set_epoch(epoch) model.train() run_loss, run_n = 0.0, 0 for img, lab, wts in dl: + if lab is None: + raise RuntimeError("training batch without labels; mixing labeled/unlabeled rows is not supported") img = img.to(device, non_blocking=True) lab = lab.to(device, non_blocking=True) wts = wts.to(device, non_blocking=True) @@ -155,6 +179,8 @@ def train_ddp(rows, val_rows, args): if rank == 0: print(f" epoch {epoch+1}/{args.epochs} loss={run_loss/max(run_n,1):.4f} " f"lr={opt.param_groups[0]['lr']:.2e}", flush=True) + 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: d, _ = evaluate(model, val_rows[: args.val_limit], device, patch[0], tta=True) print(f" [val] epoch {epoch+1} dice={d:.4f}", flush=True) @@ -163,6 +189,7 @@ def train_ddp(rows, val_rows, args): best_epoch = epoch + 1 torch.save({"model": model.state_dict(), "epoch": epoch + 1, "val_dice": best}, os.path.join(ckpt_dir, "best.pt")) + save_state() if world > 1: dist.barrier() if rank == 0: