File size: 6,238 Bytes
c3f98a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """Run the official 10-fold training matrix.
Every ``(configuration, fold)`` pair is trained with the same pinned schedule.
A run counts as complete only when it wrote ``summary.json`` - Ultralytics
emits ``weights/best.pt`` from the first epoch, so a checkpoint alone does not
mean the schedule finished. The matrix is therefore resumable across
invocations without silently keeping a half-trained fold.
Usage
-----
python scripts/03_train.py # everything still missing
python scripts/03_train.py --groups main # one group
python scripts/03_train.py --exp Exp4_2p5D_Strict --folds 0
python scripts/03_train.py --dry-run # just print the plan
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from luna_rev import config as cfg
from luna_rev import splits
from luna_rev.train import is_complete
LOG_PATH = cfg.RESULTS_DIR / "training_log.jsonl"
#: Seconds to idle after each run so Windows can reap the dataloader workers
#: before the next run starts allocating.
COOLDOWN_S = 8
def run_in_subprocess(exp, fold, force: bool = False, epochs: int | None = None) -> dict:
"""Train one pair in a fresh interpreter and return its record.
Isolating each run guarantees that workers, pinned buffers and the CUDA
context are released before the next one begins; sharing one interpreter
across 120 runs exhausted page-locked host memory partway through.
"""
import subprocess
cmd = [sys.executable, "-m", "luna_rev._train_worker", exp.name, str(fold.index)]
if force:
cmd.append("--force")
if epochs:
cmd += ["--epochs", str(epochs)]
proc = subprocess.run(cmd, cwd=str(cfg.ROOT), capture_output=True, text=True,
encoding="utf-8", errors="replace")
time.sleep(COOLDOWN_S)
for line in reversed((proc.stdout or "").splitlines()):
if line.startswith("TRAIN_RECORD "):
return json.loads(line[len("TRAIN_RECORD "):])
tail = "\n".join((proc.stderr or "").strip().splitlines()[-12:])
print(f" subprocess exit {proc.returncode}\n{tail}", flush=True)
return {"exp": exp.name, "fold": fold.index, "status": "failed",
"returncode": proc.returncode, "stderr_tail": tail}
def build_plan(groups, exp_names, fold_filter) -> list[tuple]:
"""Expand the requested selection into ``(exp, fold)`` pairs."""
folds = {f.index: f for f in splits.get_folds("official")}
experiments = cfg.experiments_for_group(*groups) if groups else list(cfg.ALL_EXPERIMENTS)
if exp_names:
wanted = set(exp_names)
experiments = [e for e in experiments if e.name in wanted]
unknown = wanted - {e.name for e in experiments}
if unknown:
raise SystemExit(f"Unknown experiment(s): {sorted(unknown)}")
plan = []
for exp in experiments:
for k in cfg.folds_for(exp):
if fold_filter is not None and k not in fold_filter:
continue
plan.append((exp, folds[k]))
return plan
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--groups", nargs="*", default=None,
help="main r_sweep w_sweep neg_ablation seeds yolo26")
ap.add_argument("--exp", nargs="*", default=None, help="explicit configuration names")
ap.add_argument("--folds", nargs="*", type=int, default=None)
ap.add_argument("--force", action="store_true", help="retrain even if weights exist")
ap.add_argument("--dry-run", action="store_true")
ap.add_argument("--epochs", type=int, default=None, help="override the schedule (benchmarking)")
args = ap.parse_args()
plan = build_plan(args.groups, args.exp, set(args.folds) if args.folds else None)
pending = [(e, f) for e, f in plan if args.force or not is_complete(e, f.index)]
print(f"planned runs : {len(plan)}")
print(f"already done : {len(plan) - len(pending)}")
print(f"to train : {len(pending)}")
if args.dry_run:
for e, f in pending:
print(f" {e.name:28s} fold{f.index} [{e.model}, r={e.r_sample}, "
f"w_min={e.w_min_px}, neg={e.negatives}, seed={e.seed}]")
return 0
if not pending:
return 0
t_start = time.time()
failures: list[tuple] = []
for i, (exp, fold) in enumerate(pending, 1):
t0 = time.time()
print(f"\n[{i}/{len(pending)}] {exp.name} fold{fold.index} "
f"(elapsed {(time.time() - t_start) / 3600:.2f} h)", flush=True)
rec = run_in_subprocess(exp, fold, force=args.force, epochs=args.epochs)
rec["wall_seconds"] = round(time.time() - t0, 1)
with LOG_PATH.open("a", encoding="utf-8") as fh:
fh.write(json.dumps(rec) + "\n")
if rec["status"] == "failed":
failures.append((exp, fold))
print(f" -> FAILED after {rec['wall_seconds'] / 60:.1f} min "
f"(continuing; will retry at the end)", flush=True)
continue
rate = (time.time() - t_start) / i
print(f" -> {rec['status']} in {rec['wall_seconds'] / 60:.1f} min "
f"| ETA {(len(pending) - i) * rate / 3600:.1f} h", flush=True)
# One retry pass: the failures we have seen are transient resource
# exhaustion, which a fresh process after an idle gap usually clears.
if failures:
print(f"\nRetrying {len(failures)} failed run(s) ...", flush=True)
still_failing = []
for exp, fold in failures:
time.sleep(30)
rec = run_in_subprocess(exp, fold, force=True, epochs=args.epochs)
with LOG_PATH.open("a", encoding="utf-8") as fh:
fh.write(json.dumps({**rec, "retry": True}) + "\n")
if rec["status"] == "failed":
still_failing.append(f"{exp.name} fold{fold.index}")
if still_failing:
print("PERMANENT FAILURES: " + ", ".join(still_failing), flush=True)
return 1
print(f"\nAll done in {(time.time() - t_start) / 3600:.2f} h")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|