Furen Xiao
a516408c4e
- Implemented 09_run_iterative.py for MONAI to orchestrate iterative pseudo-labeling rounds, including training, predicting, and evaluation. - Created 01_prepare_dataset.py, 02_plan_preprocess.py, 06_pseudo_label.py, 07_train.py, 08_eval.py, and 09_run_iterative.py for nnU-Net to manage dataset preparation, planning, training, evaluation, and iterative pseudo-labeling. - Enhanced error handling and logging throughout the scripts for better traceability. - Added functionality to handle negative pseudo-labels and TTA options during inference. - Included visualization of Dice scores across rounds for both MONAI and nnU-Net.
221 lines
No EOL
11 KiB
Python
221 lines
No EOL
11 KiB
Python
"""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() |