Buckets:
| """cwm3/train/train_overlap_head.py — F8 lane (R2 WP4 batch 2, 2026-09-17; pre-registration docs/decisions.md 2026-09-17T19:32:07Z): the detached joint | |
| driver-unmet-state × road-demand OVERLAP head (models/overlap_head.py) on the FROZEN ECAN joint model `I1_all_minus_road2drv_s<k>_final_ecan`. | |
| python train/train_overlap_head.py --seed 0 --extract-only # step 1 (GPU): trunk / marginal caches of the frozen model, both regimes, train / val / test | |
| python train/train_overlap_head.py --seed 0 --arm joint|margonly|prod # step 2/3: train the arm on the caches (GPU, seconds-minutes), predict val + test | |
| -> results/runs/OV_<arm>_s<k>/ ckpt_best.pt config.json history.jsonl TRAIN_DONE pred_overlap_{regime}_{split}.npz (val, test) PRED_DONE run_*.json | |
| Step 1 (cache, results/runs/OV_cache/, never a result): targets_{regime}_{split}.npz = the row-aligned label-free target arrays of the (regime, split) pack | |
| (S [N,25,2] hands_on / head_turned NaN-censored, S0, row_ok, M, P / Pvalid [N,25,6], P0 / P0valid, uid, t, dongle, mode; written once, shared by every seed); | |
| trunks_<frozen run>_{regime}_{split}.npz = h_u / z_e / z_vf f16 [N,128] each of the frozen model (its own forward under the pack's bf16 autocast, the model's | |
| `return_trunks`; the FINAL variant has no interaction into U or E, so these are the base-trunk states) + marg f16 [N,15,8] = the frozen model's OWN per-horizon | |
| marginals (p_hands_off, p_head_turned, pi of the six P chains: models/heads.forward_marginal of its U / P chain hazards from the anchor state, invalid anchor | |
| states filled with the TRAIN base rate) + witness: max |recomputed hazard − saved pred pack hazard| on val / test (expected <= one f16 ulp). | |
| Step 2 (train): TRAIN anchors of both regimes; AdamW 1e-3 / wd 1e-4, batch 4096, <= --epochs, early stop patience 3 on the VALIDATION masked joint NLL (the selection | |
| objective; validation only). Arms: `joint` (trunk features + anchor states + regime), `margonly` (the frozen model's own marginals + logits as per-horizon | |
| inputs, NO trunk features — attribution control), `prod` (no training: the product of the frozen model's marginals read from its prediction packs | |
| pred_{regime}_{split}.npz, cross-checked against the cache within f16 tolerance). | |
| Step 3 (predict): val / test per regime -> pred_overlap_{regime}_{split}.npz row-aligned with the causal grid (uid / t witness): q [N,15] f16 = the paper-facing | |
| P(U_{t+k} = 1, D_{t+k} = 1) (manual: U = hands_off ∨ head_turned; engaged: U = head_turned only — the engaged hand cells are out of the paper), q_hands / q_head | |
| (both readings, descriptive), J = 0.2 · sum_k q_k (+ J_hands / J_head), p_U / p_D marginals of the arm, valid = row_ok, meta json. | |
| Nothing here reads an openpilot DMS field; the CAN side enters only through the pressure targets and the frozen model's ECAN inputs. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| os.environ.setdefault("MKL_NUM_THREADS", "4") | |
| import argparse # noqa: E402 | |
| import copy # noqa: E402 | |
| import json # noqa: E402 | |
| import sys # noqa: E402 | |
| import time # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| from types import SimpleNamespace # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| for p in (CWM3, CWM3 / "data", CWM3 / "common", CWM3 / "train", CWM3 / "eval", EXPLORE / "s12_wm"): | |
| if str(p) not in sys.path: | |
| sys.path.insert(0, str(p)) | |
| import cwm3_dataset as ds # noqa: E402 | |
| from runmeta import write_run_json, code_hash # noqa: E402 | |
| RUNS = CWM3 / "results" / "runs" | |
| CACHE = RUNS / "OV_cache" | |
| KIN = CWM3 / "results" / "targets_kin_causal_routes_v2" | |
| REGIMES = ("manual", "engaged") | |
| SPLITS = ("train", "val", "test") | |
| FROZEN = "I1_all_minus_road2drv_s{seed}_final_ecan" | |
| ARMS = ("joint", "margonly", "prod") | |
| U_COLS = (4, 5) # hands_on / head_turned columns of the joint pred pack (CHAINS order V4 U2 E6) | |
| CODE_FILES = [CWM3 / "train" / "train_overlap_head.py", CWM3 / "models" / "overlap_head.py", CWM3 / "models" / "heads.py", CWM3 / "models" / "cwm3.py", | |
| CWM3 / "train" / "batch.py", CWM3 / "train" / "train_cwm3_joint.py"] | |
| def log(msg, out=None): | |
| line = f"[{time.strftime('%H:%M:%S')}] {msg}" | |
| print(line, flush=True) | |
| if out: | |
| Path(out).mkdir(parents=True, exist_ok=True) | |
| with open(Path(out) / "train.log", "a") as f: | |
| f.write(line + "\n") | |
| def u_base_rate(): | |
| """TRAIN base rate of the two anchor kin states (both regimes pooled, rows with a causal box) — the fill for invalid anchor states in the recursion.""" | |
| S0 = [] | |
| for r in REGIMES: | |
| with np.load(KIN / f"kin_causal_{r}_train.npz", allow_pickle=True) as z: | |
| ok = z["row_ok"].astype(bool) | |
| S0.append(z["S0_c"][ok].astype(np.float32)) | |
| S0 = np.concatenate(S0) | |
| return np.array([np.nanmean(S0[:, j]) for j in range(S0.shape[1])], np.float32) | |
| # ----------------------------------------------------------------------------- step 1: the frozen model's trunks + marginals | |
| def load_frozen(run_name, device, work_out): | |
| """-> ctx dict: model (eval, on device), cfg, st, pc, bank, args_data (logs into work_out), info, best.""" | |
| import train_cwm3_joint as TJ | |
| run = RUNS / run_name | |
| assert (run / "TRAIN_DONE").exists() and (run / "PRED_DONE").exists(), f"{run}: not a finished joint run" | |
| c = json.loads((run / "config.json").read_text()) | |
| cfg, a = c["cfg"], c["args"] | |
| args_model = SimpleNamespace(**a) | |
| args_model.out = str(run) # BEST / ckpt of the frozen run (read only) | |
| args_model.cpu, args_model.no_amp, args_model.synthetic, args_model.max_anchors = False, False, 0, 0 | |
| args_data = copy.copy(args_model) | |
| args_data.out = str(work_out) # every log line of the data layer goes to OUR dir, never into the frozen run | |
| ds.set_pack_version(cfg.get("pack_version", "v2")) | |
| ds.set_kin_source(cfg["kin_source"], cfg["kin_version"]) | |
| assert "road2drv" not in cfg["interactions"] and not cfg.get("naive"), "h_u / z_e must be the base-trunk states: the FINAL variant has no interaction into U or E" | |
| info = TJ.base_info_numpy(cfg, int(a["seed"])) | |
| st = {k: v for k, v in np.load(run / "stats.npz").items()} | |
| bank = TJ.load_bank(args_model, cfg, info) | |
| model, best = TJ.load_best_model(args_model, cfg, device, info) | |
| for p in model.parameters(): | |
| p.requires_grad_(False) | |
| return {"model": model, "cfg": cfg, "st": st, "pc": TJ.prep_cfgs(info), "bank": bank, "args": args_data, "info": info, "best": best, | |
| "model_id": f"{run_name}:{best}", "run": run, "TJ": TJ} | |
| def extract(ctx, regime, split, device, base_rate, max_anchors=None): | |
| """one (regime, split) -> (targets dict, trunks dict).""" | |
| import torch | |
| from models import heads as hd | |
| TJ = ctx["TJ"] | |
| args = copy.copy(ctx["args"]) | |
| args.max_anchors = int(max_anchors or 0) | |
| raw = TJ.load_raw(args, ctx["cfg"], split, subset=bool(max_anchors), regimes=[regime]) | |
| assert "P" in raw, "pressure targets missing for the pack" | |
| N = len(raw["t"]) | |
| model, st, pc = ctx["model"], ctx["st"], ctx["pc"] | |
| store = TJ.make_store(raw, device, ctx["bank"]) | |
| region_keys = model.tok_embed is not None | |
| H = {k: np.zeros((N, 128), np.float16) for k in ("h_u", "z_e", "z_vf")} | |
| marg = np.zeros((N, 15, 8), np.float16) | |
| hzU = np.zeros((N, 25, 2, 2), np.float16) | |
| hzP = np.zeros((N, 25, 6, 2), np.float16) | |
| br = torch.as_tensor(base_rate, device=device) | |
| bs = 4096 | |
| t0 = time.time() | |
| with torch.no_grad(): | |
| for i in range(0, N, bs): | |
| idx = torch.arange(i, min(i + bs, N), device=device) | |
| sl = slice(i, i + len(idx)) | |
| b, tg, _ = TJ.prep_joint(store, idx, st, pc, True, region_keys) | |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(device == "cuda")): | |
| o = model(b, return_trunks=True) | |
| tr = o["trunks"] | |
| for k in H: | |
| H[k][sl] = tr[k].float().cpu().numpy().astype(np.float16) | |
| u = hd.split_U(o["U"]) | |
| hon_u, hoff_u = hd.hazards(u["chains"]) # [B,25,2] | |
| pc_ = o["P"].float().reshape(len(idx), 25, 6, 2) | |
| hon_p, hoff_p = hd.hazards(pc_) # [B,25,6] | |
| z0u = torch.where(tg["U"]["z0valid"], tg["U"]["z0"], br[None, :].expand(len(idx), -1)) | |
| pi_u = torch.stack([hd.forward_marginal(hon_u[..., j], hoff_u[..., j], z0u[:, j]) for j in range(2)], -1)[:, :15] # [B,15,2] | |
| pi_p = torch.stack([hd.forward_marginal(hon_p[..., j], hoff_p[..., j], tg["P"]["z0"][:, j]) for j in range(6)], -1)[:, :15] | |
| marg[sl] = torch.cat([1.0 - pi_u[..., :1], pi_u[..., 1:2], pi_p], -1).cpu().numpy().astype(np.float16) | |
| hzU[sl] = torch.stack([hon_u, hoff_u], -1).cpu().numpy().astype(np.float16) | |
| hzP[sl] = torch.stack([hon_p, hoff_p], -1).cpu().numpy().astype(np.float16) | |
| if (i // bs) % 40 == 0: | |
| log(f" extract {regime}/{split} {i + len(idx):,}/{N:,} ({time.time() - t0:.0f}s)", ctx["args"].out) | |
| witness = {} | |
| pk = ctx["run"] / f"pred_{regime}_{split}.npz" | |
| if pk.exists() and not max_anchors: | |
| with np.load(pk, allow_pickle=True) as z: | |
| assert z["n_rows"] == N | |
| dU = np.abs(z["hz_on"][:, :, list(U_COLS)].astype(np.float32) - hzU[..., 0].astype(np.float32)).max() | |
| dUo = np.abs(z["hz_off"][:, :, list(U_COLS)].astype(np.float32) - hzU[..., 1].astype(np.float32)).max() | |
| dP = np.abs(z["hz_on_P"].astype(np.float32) - hzP[..., 0].astype(np.float32)).max() | |
| dPo = np.abs(z["hz_off_P"].astype(np.float32) - hzP[..., 1].astype(np.float32)).max() | |
| witness = {"pack": str(pk), "max_abs_diff_hz_on_U": float(dU), "max_abs_diff_hz_off_U": float(dUo), "max_abs_diff_hz_on_P": float(dP), "max_abs_diff_hz_off_P": float(dPo)} | |
| log(f" witness vs {pk.name}: |Δ hz| U on/off {dU:.2e}/{dUo:.2e}, P on/off {dP:.2e}/{dPo:.2e}", ctx["args"].out) | |
| tgt = {"uid": raw["uid"], "t": raw["t"], "dongle": raw["dongle"], "mode": raw["mode"].astype(np.int8), "S": raw["S"].astype(np.float16), "S0": raw["S0"].astype(np.float32), | |
| "row_ok": raw["kin_row_ok"].astype(bool), "M": raw["M"].astype(bool), "P": raw["P"].astype(bool), "Pvalid": raw["Pvalid"].astype(bool), | |
| "P0": raw["P0"].astype(bool), "P0valid": raw["P0valid"].astype(bool), "idx": raw["idx"]} | |
| trk = dict(H, marg=marg, uid=raw["uid"], t=raw["t"], model_id=np.array(ctx["model_id"]), witness=np.array(json.dumps(witness)), base_rate=base_rate) | |
| del store, raw | |
| if device == "cuda": | |
| torch.cuda.empty_cache() | |
| return tgt, trk | |
| def cache_paths(frozen, regime, split): | |
| return CACHE / f"targets_{regime}_{split}.npz", CACHE / f"trunks_{frozen}_{regime}_{split}.npz" | |
| def ensure_cache(frozen, regime, split, device, ctx_holder, work_out, max_anchors=None): | |
| """-> (targets dict, trunks dict); builds the missing cache files once (smoke subsets are never cached).""" | |
| pt, pz = cache_paths(frozen, regime, split) | |
| if not max_anchors and pt.exists() and pz.exists(): | |
| with np.load(pt, allow_pickle=True) as z: | |
| tgt = {k: z[k] for k in z.files} | |
| with np.load(pz, allow_pickle=True) as z: | |
| trk = {k: z[k] for k in z.files} | |
| assert np.array_equal(tgt["uid"], trk["uid"]), "target / trunk cache row mismatch" | |
| return tgt, trk | |
| if ctx_holder.get("ctx") is None: | |
| ctx_holder["ctx"] = load_frozen(frozen, device, work_out) | |
| ctx_holder["base_rate"] = u_base_rate() | |
| log(f"frozen {ctx_holder['ctx']['model_id']} loaded; anchor base rate {ctx_holder['base_rate'].tolist()}", work_out) | |
| t0 = time.time() | |
| tgt, trk = extract(ctx_holder["ctx"], regime, split, device, ctx_holder["base_rate"], max_anchors) | |
| log(f"extracted {regime}/{split}: {len(tgt['t']):,} rows in {time.time() - t0:.0f}s", work_out) | |
| if not max_anchors: | |
| CACHE.mkdir(parents=True, exist_ok=True) | |
| if pt.exists(): | |
| with np.load(pt, allow_pickle=True) as z: | |
| assert np.array_equal(z["uid"], tgt["uid"]) and np.allclose(z["t"], tgt["t"]), "existing target cache differs" | |
| else: | |
| np.savez(pt, **tgt) | |
| np.savez(pz, **trk) | |
| return tgt, trk | |
| # ----------------------------------------------------------------------------- tensors | |
| def anchor_features(tgt): | |
| S0 = tgt["S0"].astype(np.float32) | |
| v = np.isfinite(S0) | |
| A = np.concatenate([np.nan_to_num(S0), v.astype(np.float32), tgt["P0"].astype(np.float32) * tgt["P0valid"], tgt["P0valid"].astype(np.float32)], 1) | |
| return A.astype(np.float32) | |
| def zk_features(trk): | |
| """margonly inputs: the frozen model's own marginals and their logits, [N,15,16].""" | |
| m = np.clip(trk["marg"].astype(np.float32), 1e-4, 1 - 1e-4) | |
| return np.concatenate([m, np.log(m / (1 - m))], -1).astype(np.float16) | |
| def tensors(parts, arm, device): | |
| """list of (tgt, trk) -> dict of tensors on device.""" | |
| import torch | |
| from models.overlap_head import build_targets | |
| Z = np.concatenate([np.concatenate([tr["h_u"], tr["z_e"], tr["z_vf"]], 1) for _, tr in parts]) if arm == "joint" else None | |
| A = np.concatenate([anchor_features(tg) for tg, _ in parts]) | |
| mode = np.concatenate([tg["mode"] for tg, _ in parts]).astype(np.int64) | |
| ys, obs = [], [] | |
| for tg, _ in parts: | |
| y3, o3 = build_targets(tg["S"], tg["P"], tg["Pvalid"], tg["M"], tg["row_ok"]) | |
| ys.append(y3); obs.append(o3) | |
| out = {"A": torch.from_numpy(A).to(device), "mode": torch.from_numpy(mode).to(device), "y3": torch.from_numpy(np.concatenate(ys)).to(device), | |
| "obs3": torch.from_numpy(np.concatenate(obs)).to(device), "N": len(mode)} | |
| out["Z"] = torch.from_numpy(Z).to(device) if Z is not None else None | |
| out["zk"] = torch.from_numpy(np.concatenate([zk_features(tr) for _, tr in parts])).to(device) if arm == "margonly" else None | |
| return out | |
| def fwd(model, T, idx): | |
| import torch | |
| z = T["Z"][idx] if T["Z"] is not None else torch.zeros((len(idx), 1), device=T["A"].device) | |
| return model(z, T["A"][idx], T["mode"][idx], None if T["zk"] is None else T["zk"][idx]) | |
| def run_epoch(model, opt, T, bs, train, gen=None): | |
| import torch | |
| from models.overlap_head import overlap_loss | |
| idx_all = torch.arange(T["N"], device=T["A"].device) | |
| order = idx_all[torch.randperm(T["N"], generator=gen, device=idx_all.device)] if train else idx_all | |
| tot, n = 0.0, 0.0 | |
| for i in range(0, T["N"], bs): | |
| idx = order[i:i + bs] | |
| logits = fwd(model, T, idx) | |
| loss, parts = overlap_loss(logits, T["y3"][idx], T["obs3"][idx]) | |
| if train: | |
| opt.zero_grad(set_to_none=True) | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) | |
| opt.step() | |
| w = float(parts["n_obs"]) | |
| tot += float(loss.detach()) * w | |
| n += w | |
| return tot / max(n, 1.0) | |
| # ----------------------------------------------------------------------------- prediction files | |
| def write_pred(out, regime, split, tgt, q_hands, q_head, pU, pD, meta): | |
| from models.overlap_head import DT | |
| use_hands = regime == "manual" | |
| q = q_hands if use_hands else q_head | |
| m = dict(meta, regime=regime, split=split, n_rows=int(len(tgt["t"])), dt=DT, horizon=15, w="uniform (w_k = 1)", | |
| paper_facing_U="hands_off | head_turned" if use_hands else "head_turned only (engaged hand cells out of the paper)", | |
| written_utc=time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())) | |
| np.savez_compressed(Path(out) / f"pred_overlap_{regime}_{split}.npz", uid=tgt["uid"], t=tgt["t"], dongle=tgt["dongle"], valid=tgt["row_ok"].astype(bool), | |
| q=q.astype(np.float16), q_hands=q_hands.astype(np.float16), q_head=q_head.astype(np.float16), | |
| J=(DT * q.astype(np.float32).sum(1)).astype(np.float32), J_hands=(DT * q_hands.astype(np.float32).sum(1)).astype(np.float32), | |
| J_head=(DT * q_head.astype(np.float32).sum(1)).astype(np.float32), p_U=pU.astype(np.float16), p_D=pD.astype(np.float16), | |
| meta=np.array(json.dumps(m, default=str))) | |
| def predict_arm(model, T, tgt, out, regime, split, device, meta, bs=16384): | |
| import torch | |
| from models.overlap_head import overlap_from_joint, marginals_from_joint | |
| N = T["N"] | |
| qh, qd, pU, pD = (np.zeros((N, 15), np.float32) for _ in range(4)) | |
| model.eval() | |
| with torch.no_grad(): | |
| for i in range(0, N, bs): | |
| idx = torch.arange(i, min(i + bs, N), device=device) | |
| P8 = torch.softmax(fwd(model, T, idx).float(), -1) | |
| qh[i:i + len(idx)] = overlap_from_joint(P8, True).cpu().numpy() | |
| qd[i:i + len(idx)] = overlap_from_joint(P8, False).cpu().numpy() | |
| mg = marginals_from_joint(P8) | |
| pU[i:i + len(idx)] = mg["p_U"].cpu().numpy() | |
| pD[i:i + len(idx)] = mg["p_D"].cpu().numpy() | |
| write_pred(out, regime, split, tgt, qh, qd, pU, pD, meta) | |
| return N | |
| def product_from_pack(run, regime, split, base_rate, rows=None): | |
| """the pre-registered product baseline from the frozen run's own prediction pack -> (p_hands_off, p_head, pi_P [N,15,6]) float32; rows = pack rows (smoke subsets).""" | |
| import torch | |
| from models import heads as hd | |
| with np.load(run / f"pred_{regime}_{split}.npz", allow_pickle=True) as z: | |
| r = np.arange(int(z["n_rows"])) if rows is None else np.asarray(rows, np.int64) | |
| hon, hoff, z0, z0v = z["hz_on"][r][:, :, list(U_COLS)].astype(np.float32), z["hz_off"][r][:, :, list(U_COLS)].astype(np.float32), z["z0"][r][:, list(U_COLS)].astype(np.float32), z["z0valid"][r][:, list(U_COLS)] | |
| honP, hoffP, z0P = z["hz_on_P"][r].astype(np.float32), z["hz_off_P"][r].astype(np.float32), z["z0_P"][r].astype(np.float32) | |
| z0f = np.where(z0v, z0, base_rate[None, :2]) | |
| with torch.no_grad(): | |
| pu = torch.stack([hd.forward_marginal(torch.from_numpy(hon[..., j]), torch.from_numpy(hoff[..., j]), torch.from_numpy(z0f[:, j])) for j in range(2)], -1)[:, :15].numpy() | |
| pp = torch.stack([hd.forward_marginal(torch.from_numpy(honP[..., j]), torch.from_numpy(hoffP[..., j]), torch.from_numpy(z0P[:, j])) for j in range(6)], -1)[:, :15].numpy() | |
| return 1.0 - pu[..., 0], pu[..., 1], pp | |
| def predict_prod(frozen, tgt, trk, out, regime, split, base_rate, meta): | |
| import torch | |
| from models.overlap_head import joint_from_marginals, overlap_from_joint, marginals_from_joint | |
| p_ho, p_ht, pi_p = product_from_pack(RUNS / frozen, regime, split, base_rate, tgt["idx"]) | |
| cache_m = trk["marg"].astype(np.float32) | |
| dmax = float(max(np.abs(cache_m[..., 0] - p_ho).max(), np.abs(cache_m[..., 1] - p_ht).max(), np.abs(cache_m[..., 2:] - pi_p).max())) | |
| pD = 1.0 - np.prod(1.0 - pi_p, -1) | |
| P8 = joint_from_marginals(torch.from_numpy(p_ho), torch.from_numpy(p_ht), torch.from_numpy(pD)) | |
| qh, qd = overlap_from_joint(P8, True).float().numpy(), overlap_from_joint(P8, False).float().numpy() | |
| mg = marginals_from_joint(P8) | |
| write_pred(out, regime, split, tgt, qh, qd, mg["p_U"].float().numpy(), mg["p_D"].float().numpy(), dict(meta, pack_vs_cache_max_abs_diff=dmax)) | |
| return len(tgt["t"]), dmax | |
| # ----------------------------------------------------------------------------- main | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--frozen", default="", help=f"frozen joint run (default {FROZEN})") | |
| ap.add_argument("--arm", default="joint", choices=list(ARMS)) | |
| ap.add_argument("--extract-only", action="store_true", help="build the caches (both regimes, train / val / test) and stop") | |
| ap.add_argument("--epochs", type=int, default=20) | |
| ap.add_argument("--patience", type=int, default=3) | |
| ap.add_argument("--batch", type=int, default=4096) | |
| ap.add_argument("--lr", type=float, default=1e-3) | |
| ap.add_argument("--d", type=int, default=96) | |
| ap.add_argument("--max-anchors", type=int, default=0, help="smoke subset per (regime, split); never cached, never a result (out dir _smoke)") | |
| ap.add_argument("--cpu", action="store_true") | |
| ap.add_argument("--out", default="") | |
| args = ap.parse_args() | |
| import torch | |
| from models.overlap_head import OverlapHead, n_params, D_ANCHOR, K | |
| torch.set_num_threads(4) | |
| device = "cuda" if torch.cuda.is_available() and not args.cpu else "cpu" | |
| frozen = args.frozen or FROZEN.format(seed=args.seed) | |
| smoke = "_smoke" if args.max_anchors else "" | |
| out = Path(args.out) if args.out else RUNS / (f"OV_extract_s{args.seed}{smoke}" if args.extract_only else f"OV_{args.arm}_s{args.seed}{smoke}") | |
| out.mkdir(parents=True, exist_ok=True) | |
| t_start = time.time() | |
| log(f"=== train_overlap_head seed {args.seed} arm {args.arm} frozen {frozen} extract_only {args.extract_only} out {out} device {device} ===", out) | |
| holder = {"ctx": None} | |
| mx = args.max_anchors or None | |
| data = {} | |
| for split in SPLITS: | |
| if args.arm == "prod" and split == "train" and not args.extract_only: | |
| continue # the product baseline needs no training rows | |
| data[split] = [ensure_cache(frozen, r, split, device, holder, out, mx) for r in REGIMES] | |
| gpu_extract_s = round(time.time() - t_start, 1) | |
| meta = {"frozen_run": frozen, "model_id": str(data["val"][0][1]["model_id"]), "arm": args.arm, "seed": args.seed, "smoke": bool(mx), | |
| "witness": {f"{r}_{s}": json.loads(str(data[s][i][1]["witness"])) for s in data for i, r in enumerate(REGIMES)}} | |
| if holder["ctx"] is not None: | |
| del holder["ctx"] | |
| torch.cuda.empty_cache() if device == "cuda" else None | |
| if args.extract_only: | |
| (out / "EXTRACT_DONE").write_text(json.dumps({"gpu_wall_s": gpu_extract_s, "witness": meta["witness"], "cache": str(CACHE)}, indent=1)) | |
| write_run_json(str(out), f"overlap_extract_s{args.seed}", CODE_FILES, [], seed=args.seed, t_start=t_start, | |
| extra={"frozen_run": frozen, "witness": meta["witness"], "gpu_wall_s": gpu_extract_s, "dms_fields_used": []}) | |
| log(f"caches built in {gpu_extract_s:.0f}s; --extract-only", out) | |
| return | |
| base_rate = np.asarray(data["val"][0][1]["base_rate"], np.float32) | |
| n_pred, npar, hist, best, best_ep = {}, 0, [], float("nan"), -1 | |
| if args.arm == "prod": | |
| (out / "TRAIN_DONE").write_text(json.dumps({"no_train": True, "params": 0, "selection": "none (product of the frozen model's marginals)"})) | |
| (out / "BEST").write_text("none\n") | |
| for split in ("val", "test"): | |
| for (tgt, trk), r in zip(data[split], REGIMES): | |
| n, dmax = predict_prod(frozen, tgt, trk, out, r, split, base_rate, meta) | |
| n_pred[f"{r}_{split}"] = n | |
| log(f"pred_overlap_{r}_{split}.npz written ({n:,} rows; pack vs cache marginals max |Δ| {dmax:.2e})", out) | |
| else: | |
| Ttr, Tva = tensors(data["train"], args.arm, device), tensors(data["val"], args.arm, device) | |
| log(f"train {Ttr['N']:,} rows, val {Tva['N']:,} rows; observed (row, k) share train {float(Ttr['obs3'].any(-1).float().mean()):.4f}", out) | |
| torch.manual_seed(args.seed) | |
| model = OverlapHead(d_z=384, d_anchor=D_ANCHOR, d=args.d, n_h=K, d_k=(16 if args.arm == "margonly" else 0), use_z=(args.arm == "joint")).to(device) | |
| npar = n_params(model) | |
| log(f"OverlapHead[{args.arm}] params {npar:,}", out) | |
| opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) | |
| gen = torch.Generator(device=device).manual_seed(args.seed) | |
| best, bad = float("inf"), 0 | |
| for ep in range(1, args.epochs + 1): | |
| t0 = time.time() | |
| model.train() | |
| ltr = run_epoch(model, opt, Ttr, args.batch, True, gen) | |
| model.eval() | |
| with torch.no_grad(): | |
| lva = run_epoch(model, None, Tva, 16384, False) | |
| rec = {"epoch": ep, "train_joint_nll": ltr, "val_joint_nll": lva, "wall_s": round(time.time() - t0, 1)} | |
| hist.append(rec) | |
| with open(out / "history.jsonl", "a") as f: | |
| f.write(json.dumps(rec) + "\n") | |
| improved = lva < best - 1e-6 | |
| if improved: | |
| best, best_ep, bad = lva, ep, 0 | |
| torch.save({"model": model.state_dict(), "epoch": ep, "val_joint_nll": lva, "args": vars(args), "frozen": frozen, "arm": args.arm, "d": args.d, | |
| "base_rate": base_rate}, out / "ckpt_best.pt") | |
| else: | |
| bad += 1 | |
| log(f"ep {ep} train {ltr:.5f} val {lva:.5f} {'*' if improved else ''} ({rec['wall_s']}s)", out) | |
| if bad >= args.patience: | |
| break | |
| z = torch.load(out / "ckpt_best.pt", map_location=device, weights_only=False) | |
| model.load_state_dict(z["model"]) | |
| (out / "BEST").write_text("ckpt_best.pt\n") | |
| (out / "TRAIN_DONE").write_text(json.dumps({"best_epoch": best_ep, "best": best, "epochs_run": len(hist), "params": npar, "selection": "validation masked joint NLL"})) | |
| del Ttr | |
| for split in ("val", "test"): | |
| T = Tva if split == "val" else tensors(data["test"], args.arm, device) | |
| off = 0 | |
| for (tgt, trk), r in zip(data[split], REGIMES): | |
| n = len(tgt["t"]) | |
| Tsub = {k: (v[off:off + n] if torch.is_tensor(v) else v) for k, v in T.items()} | |
| Tsub["N"] = n | |
| n_pred[f"{r}_{split}"] = predict_arm(model, Tsub, tgt, out, r, split, device, dict(meta, params=npar, best_epoch=best_ep, val_joint_nll=best)) | |
| off += n | |
| log(f"pred_overlap_{r}_{split}.npz written ({n:,} rows)", out) | |
| (out / "config.json").write_text(json.dumps({"args": vars(args), "frozen": frozen, "arm": args.arm, "params": npar, "best_epoch": best_ep, "val_joint_nll": best, | |
| "meta": meta, "code_sha256": code_hash(CODE_FILES)}, indent=1, default=str)) | |
| inputs = [] if mx else sorted({RUNS / frozen / (RUNS / frozen / "BEST").read_text().strip()} | {cache_paths(frozen, r, s)[0] for r in REGIMES for s in ("val", "test")}) | |
| write_run_json(str(out), f"train_overlap_head_{args.arm}_s{args.seed}", CODE_FILES, inputs, seed=args.seed, t_start=t_start, | |
| extra={"frozen_run": frozen, "arm": args.arm, "params": npar, "best_epoch": best_ep, "val_joint_nll": best, "history": hist, "n_pred": n_pred, | |
| "selection": "validation masked joint NLL", "smoke": bool(mx), "dms_fields_used": [], "gpu_wall_s": round(time.time() - t_start, 1), | |
| "gpu_extract_s_in_this_process": gpu_extract_s}) | |
| (out / "PRED_DONE").write_text(json.dumps({"splits": ["val", "test"], "regimes": list(REGIMES), "n_pred": n_pred})) | |
| log(f"DONE in {time.time() - t_start:.0f}s", out) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 27.6 kB
- Xet hash:
- ab4e5c65fce6480072121f0eba406e5f521da1f3dbcdbc927f06a40b52a8c1cb
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.