xfr/longitudinal
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(train): implement checkpointing and robust data loading

Browse source 
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.
This commit is contained in:
Furen Xiao 2026-09-26 03:44:12 +08:00
parent 33bc3c603f
commit a491ae0523
3 changed files with 84 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
scripts/test_dataloader.py Normal file
View file

@ -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()

48
src/dataset.py
View file

@ -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

33
src/training.py
View file

@ -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: