"""MONAI DDP training entrypoint for one pseudo-labeling round (torchrun). 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 --batch 3 --ckpt-dir runs_monai/round0 \ [--pretrained runs_monai/round0/best.pt] 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); every save carries a "b_aligned" marker that 03/04 read for the matching Z-score input protocol """ import argparse import os import sys 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__)))) 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, ckpt_b_aligned) def main(): ap = argparse.ArgumentParser() 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, 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 (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() torch.manual_seed(0) np.random.seed(0) torch.cuda.set_device(local_rank) device = f"cuda:{local_rank}" 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, b_aligned=ba) steps_per_epoch = max(len(dl), 1) 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 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"]) 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) barrier() def save_state(): tmp = state_f + ".tmp" 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) for epoch in range(start_epoch, args.epochs): 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) lab = batch["plabel"].squeeze(1).to(device, non_blocking=True) wts = batch["w"].to(device, non_blocking=True) with torch.autocast("cuda", dtype=torch.bfloat16): logits = ddp(img) 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(), 12.0 if ba else 5.0) opt.step() if sched is not None: sched.step() run_loss += float(loss.detach()) run_n += 1 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: # 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, 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 = 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, "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, "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() if __name__ == "__main__": main()