133 lines
No EOL
5.7 KiB
Python
133 lines
No EOL
5.7 KiB
Python
"""Orchestrates the iterative pseudo-labeling study.
|
|
|
|
Round 0: baseline supervised training on labeled T1c (patient-level split).
|
|
Round k: pseudo-label the unlabeled pool with round k-1 model (confidence +
|
|
volume plausibility + longitudinal consistency), then fine-tune with
|
|
labeled + pseudo data (warm restart, lower LR). After each round the
|
|
model is evaluated on the held-out patient-level test split.
|
|
|
|
Usage: python scripts/09_run_iterative.py [--rounds 4] [--gpus 3]
|
|
"""
|
|
import os
|
|
import sys
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
import argparse
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
from src.common import ROOT, DATA, load_jsonl, save_jsonl
|
|
|
|
|
|
def run(cmd, log, retries=2):
|
|
for attempt in range(retries + 1):
|
|
try:
|
|
with open(log, "a") as f:
|
|
f.write(f"$ (attempt {attempt + 1}) " + cmd + "\n")
|
|
print(f"$ (attempt {attempt + 1}) " + cmd, flush=True)
|
|
subprocess.run(cmd, shell=True, cwd=ROOT, stdout=f, stderr=subprocess.STDOUT, check=True)
|
|
return
|
|
except subprocess.CalledProcessError:
|
|
if attempt == retries:
|
|
raise
|
|
print(f"[orch] command failed, retrying in 60s: {cmd}", flush=True)
|
|
import time
|
|
time.sleep(60)
|
|
|
|
|
|
def torchrun(nproc, script, extra):
|
|
tr = shutil.which("torchrun") or (sys.executable + " -m torch.distributed.run")
|
|
return f"{tr} --standalone --nproc_per_node {nproc} {script} {extra}"
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--rounds", type=int, default=4)
|
|
ap.add_argument("--gpus", type=int, default=3)
|
|
ap.add_argument("--base-epochs", type=int, default=40)
|
|
ap.add_argument("--round-epochs", type=int, default=12)
|
|
ap.add_argument("--lr", type=float, default=3e-4)
|
|
ap.add_argument("--round-lr", type=float, default=8e-5)
|
|
ap.add_argument("--pseudo-weight", type=float, default=0.3)
|
|
ap.add_argument("--batch", type=int, default=3)
|
|
args = ap.parse_args()
|
|
|
|
man = os.path.join(DATA, "manifests")
|
|
train_f = os.path.join(man, "split_train.jsonl")
|
|
val_f = os.path.join(man, "split_val.jsonl")
|
|
test_f = os.path.join(man, "split_test.jsonl")
|
|
pool_f = os.path.join(man, "unlabeled_pool.jsonl")
|
|
results_dir = f"{ROOT}/results"
|
|
os.makedirs(results_dir, exist_ok=True)
|
|
logdir = f"{ROOT}/logs"
|
|
os.makedirs(logdir, exist_ok=True)
|
|
main_log = os.path.join(logdir, "iterative.log")
|
|
|
|
# ---- round 0: baseline ----
|
|
run(torchrun(args.gpus, "scripts/07_train.py",
|
|
f"--rows {train_f} --val {val_f} --epochs {args.base_epochs} --lr {args.lr} "
|
|
f"--batch {args.batch} --ckpt-dir runs/round0"), f"{main_log}")
|
|
run(f"python scripts/08_eval.py --rows {test_f} --ckpt runs/round0/best.pt", f"{main_log}")
|
|
|
|
for k in range(1, args.rounds + 1):
|
|
pdir = os.path.join(DATA, "pseudo/round" + str(k))
|
|
run(torchrun(args.gpus, "scripts/06_pseudo_label.py",
|
|
f"--ckpt runs/round{k-1}/best.pt --unlabeled {pool_f} --out {pdir}"), f"{main_log}")
|
|
# build round-k training manifest: labeled + accepted pseudo rows (weighted)
|
|
pr = load_jsonl(os.path.join(pdir, "rows.jsonl"))
|
|
used = []
|
|
for r in pr:
|
|
if r["role"] == "pos":
|
|
r["w"] = args.pseudo_weight
|
|
used.append(r)
|
|
elif r["role"] == "neg":
|
|
r["w"] = args.pseudo_weight
|
|
r["is_neg"] = True
|
|
r["plabel"] = None
|
|
used.append(r)
|
|
trf = os.path.join(pdir, "train_rows.jsonl")
|
|
save_jsonl(used, trf)
|
|
run(torchrun(args.gpus, "scripts/07_train.py",
|
|
f"--rows {train_f},{trf} --val {val_f} --epochs {args.round_epochs} "
|
|
f"--lr {args.round_lr} --batch {args.batch} --resume runs/round{k-1}/best.pt "
|
|
f"--ckpt-dir runs/round{k} --val-every 1"), f"{main_log}")
|
|
run(f"python scripts/08_eval.py --rows {test_f} --ckpt runs/round{k}/best.pt", f"{main_log}")
|
|
|
|
# ---- report ----
|
|
table = []
|
|
for k in list(range(args.rounds + 1)):
|
|
f = os.path.join(results_dir, f"round{k}_test.json")
|
|
if os.path.exists(f):
|
|
r = json.load(open(f))
|
|
table.append({"round": k, "test_dice": round(r["dice"], 4), "n_test": r["n"],
|
|
"ckpt_epoch": r.get("epoch")})
|
|
pf = os.path.join(DATA, f"pseudo/round{k}/summary.json") if k > 0 else None
|
|
if pf and os.path.exists(pf):
|
|
s = json.load(open(pf))
|
|
table[-1].update({"n_pos": s["n_pos"], "n_neg": s["n_neg"],
|
|
"n_rej_cons": s["n_rejected_consistency"],
|
|
"pos_vol_med_mm3": s["pos_vol_mm3"]["med"]})
|
|
sf = os.path.join(ROOT, f"runs/round{k}/final.pt")
|
|
table[-1]["ckpt"] = sf if os.path.exists(sf) else ""
|
|
save_jsonl(table, os.path.join(results_dir, "iterative_table.jsonl"))
|
|
print(json.dumps(table, indent=1))
|
|
try:
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
rs = [t["round"] for t in table]
|
|
ds = [t["test_dice"] for t in table]
|
|
plt.figure(figsize=(6, 4))
|
|
plt.plot(rs, ds, "o-")
|
|
plt.xlabel("pseudo-labeling round")
|
|
plt.ylabel("holdout tumor Dice")
|
|
for x, y in zip(rs, ds):
|
|
plt.annotate(f"{y:.3f}", (x, y), textcoords="offset points", xytext=(0, 8), fontsize=8)
|
|
plt.grid(alpha=0.3)
|
|
plt.tight_layout()
|
|
plt.savefig(os.path.join(results_dir, "iterative_dice.png"), dpi=150)
|
|
except Exception as e: # noqa
|
|
print("plot failed:", repr(e))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |