Buckets:
| """Generic trainer for keypoint-only forecasters on the AIDE protocol. | |
| python aide/train/train_kp.py --model simlpe_kp --seed 42 --lr 3e-4 --epochs 300 --patience 60 | |
| python aide/train/train_kp.py --model cwm3d --seed 42 # plan §4 recipe defaults | |
| * data: cache S [N,45,136,3] (train windows) / S10 (val + test, protocol frames) | |
| * train windows (--aug-window): with p=0.5 the protocol grid sample_idx(45); otherwise a random | |
| 10-index grid whose 9 gaps are in {4,5} (the protocol grid is a member of this family); grids are | |
| drawn vectorised per batch and the 45-frame windows live on the GPU (no per-sample Python loop) | |
| * CPU threads capped at 4 (OMP/MKL env + torch.set_num_threads) -- the host is shared | |
| * augmentation (train only): keypoint noise σ·(1−conf) px [--noise-px], joint dropout [--joint-drop] | |
| (conf->0, xy->last-observed xy of that joint), translate ±[--translate] px and scale ×[1±--scale] | |
| about the last-observed centroid (applied to all 10 frames), horizontal flip [--flip-p] | |
| * models: simlpe_kp* (aide/models/simlpe_kp.py, tensor output) and cwm3d* (aide/models/cwm3d.py, | |
| dict output with log_b) -- losses via aide/models/losses.py (--w-euc/--w-nll/--w-bone/--w-jerk); | |
| driverwm_repro (aide/models/driverwm_repro.py, features-only input, its own paper-faithful loss); | |
| driverwm_kp (aide/models/driverwm_kp.py, arch attribution (2): Driver-WM dual stream with keypoint part tokens in the internal | |
| stream + exterior Qwen features, tensor output, shared recipe; the exterior features are loaded as for --visual) | |
| * visual tokens (--visual [--vis-key pooled|vis_pooled] [--traffic-gate]): pooled Qwen3-VL features from | |
| results/feats/feats_qwen.npz; a training sub-window frame takes the feature of the NEAREST protocol | |
| frame (features exist for the 10 protocol frames only; max offset 2 frames); the model sees the 5 | |
| observed frames' features only. Semantic heads (--semantic): CE with label smoothing and | |
| class-balanced weights (train frequencies), lambda_sem per head; logits_<task> saved in pred npz. | |
| * AdamW + cosine (linear warm-up), EMA weights, early stopping on val All-MPJPE (EMA model); with --dual-select | |
| --patience-f1 K the stop waits until val mean macro-F1 has also stalled for K epochs (roadmap B, A1) | |
| * --motion-oversample k: TRAIN clips drawn with replacement per epoch, p ∝ 1 + k·(GT future displacement / mean) (B1) | |
| * baselines on the same recipe: simlpe_kp, motionbert_kp (aide/models/motionbert_kp.py), stgcn_kp (aide/models/stgcn_kp.py) | |
| * --formulation {residual,abs} (docs/decisions.md 2026-09-15, explanatory ablation): `residual` (default) = unchanged behaviour; | |
| `abs` = motionbert_kp* / cwm3d* read out ABSOLUTE normalised coordinates (pred = out ⊙ (W, H) px, default non-zero init, no | |
| residual to the last pose, no anchors); every other model rejects `abs` | |
| * HM h=5 arms (docs/decisions.md 2026-09-17T07:00:22Z, part (A); defaults = unchanged behaviour, aide/tests/test_hm_h5_arms.py): | |
| --hand-weight w (arms a1/a2: the 42 hand keypoints weighted w in the Euclidean term, other joints 1, metrics unchanged), | |
| --horizon-gain learned|parts (arms b/c, cwm3d* only: learned per-horizon gain on the anchored residual, heads.HorizonGain), | |
| --horizon-gain-parts (alias of --horizon-gain parts); the learned gains are written to the run json (`hgain`) | |
| * F7 ensemble distillation (docs/decisions.md 2026-09-17 F7 pre-registration; defaults = unchanged behaviour, aide/tests/test_f7_distil_pretrain.py): | |
| --distil-teachers run1,run2,... (run dirs under results/runs with ckpt_best.pt of the SAME architecture; loaded strict, eval, no grad) and | |
| --w-distil w (coordinates: Euclidean between the student prediction and the MEAN teacher prediction on the same augmented observed window), | |
| --w-distil-modes w (cwm3d anchor modes: KL(mean teacher anchor_w || student anchor_w)), --w-distil-logb w (L1 on log_b vs the mean teacher | |
| log_b); the GT loss is unchanged; every weight 0 (default) = no teacher is loaded and the code path is the old one | |
| * GPU work under fcntl lock on /tmp/s13_gpu.lock (--no-lock inside a batch wrapper); --max-minutes | |
| * outputs results/runs/<name>/{pred_val.npz,pred_test.npz,ckpt_best.pt,history.json}, harness rows | |
| (val + test, with paired Δ vs ZV and vs --refs) and run_<name>_<utc>.json. | |
| * --defer-eval: train + predict only (writes pending_eval.json); `finish_eval(run_dir)` (or | |
| `--eval-only --name <run>`) runs the harness + run json later -- lets a batch driver hold the GPU | |
| lock for several trainings and do the CPU bootstrap after releasing it. | |
| * --out-root / --no-csv: redirect outputs (smoke tests) without touching results/. | |
| Leak guard: the model only receives obs = frames[:T_OBS]; targets are frames[T_OBS:]. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| # CPU thread cap (host is shared with decode-bound extraction jobs): must precede numpy / torch import | |
| for _v in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS"): | |
| os.environ.setdefault(_v, "4") | |
| import argparse # noqa: E402 | |
| import copy # noqa: E402 | |
| import fcntl # noqa: E402 | |
| import json # noqa: E402 | |
| import math # noqa: E402 | |
| import random # noqa: E402 | |
| import sys # noqa: E402 | |
| import time # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| torch.set_num_threads(int(os.environ["OMP_NUM_THREADS"])) | |
| ROOT = Path(__file__).resolve().parents[2] | |
| sys.path.insert(0, str(ROOT)) | |
| from aide.data import protocol as P # noqa: E402 | |
| from aide.eval import metrics as M # noqa: E402 | |
| from aide.eval.harness import parse_refs, run_harness, summary_line # noqa: E402 | |
| from aide.models import losses as LS # noqa: E402 | |
| import torch.nn.functional as TF # noqa: E402 | |
| from common.runmeta import write_run_json # noqa: E402 | |
| GPU_LOCK = "/tmp/s13_gpu.lock" | |
| DEFAULT_REFS = ["zvshrink=zv_shrink_val", "simlpe=simlpe_kp_s42_v2"] | |
| FEATS_NPZ = P.RESULTS / "feats" / "feats_qwen.npz" | |
| EGO_NPZ = P.RESULTS / "feats" / "feats_egomotion.npz" | |
| PROTO_IDX = P.sample_idx(P.N_FRAMES_STD) | |
| # ----------------------------------------------------------------------------- models | |
| def _formulation_kw(name: str, a, supported: bool) -> dict: | |
| """--formulation (decisions 2026-09-15): {} for the default `residual` (build calls unchanged), {"formulation": "abs"} for the two | |
| supported families (motionbert_kp*, cwm3d* built by aide/models/cwm3d.py); ValueError for every other model.""" | |
| form = getattr(a, "formulation", "residual") | |
| if form == "residual": | |
| return {} | |
| if not supported: | |
| raise ValueError(f"--formulation {form} is implemented for motionbert_kp* and cwm3d* (aide/models/cwm3d.py) only, not for {name}") | |
| return {"formulation": form} | |
| def build_model(name: str, a) -> torch.nn.Module: | |
| hg = getattr(a, "horizon_gain", "none") | |
| if hg not in (None, "none") and not (name.startswith("cwm3d") and name not in ("cwm3d_v2", "cwm3d_kp_spatial", "cwm3d_plain")): | |
| raise ValueError(f"--horizon-gain {hg} is implemented for cwm3d* (aide/models/cwm3d.py) only, not for {name}") | |
| if name.startswith("simlpe_kp"): | |
| _formulation_kw(name, a, False) | |
| from aide.models import simlpe_kp | |
| return simlpe_kp.build(name, dropout=a.dropout), [simlpe_kp.__file__] | |
| if name.startswith("motionbert_kp"): | |
| from aide.models import motionbert_kp | |
| return motionbert_kp.build(name, dropout=a.dropout, **_formulation_kw(name, a, True)), [motionbert_kp.__file__] | |
| if name.startswith("stgcn_kp"): | |
| _formulation_kw(name, a, False) | |
| from aide.models import driverwm_repro, motionbert_kp, stgcn_kp | |
| return stgcn_kp.build(name, dropout=a.dropout), [stgcn_kp.__file__, driverwm_repro.__file__, motionbert_kp.__file__] | |
| if name == "sem_only": | |
| _formulation_kw(name, a, False) | |
| from aide.models import cwm3d, heads | |
| return cwm3d.build("sem_only", use_ego=bool(getattr(a, "egomotion", False)), sem_view_gate=bool(getattr(a, "sem_view_gate", False)), | |
| sem_frame_avg=tuple(x for x in getattr(a, "sem_frame_avg", "").split(",") if x)), [cwm3d.__file__, heads.__file__] | |
| if name == "cwm3d_v2": | |
| # K-v2 (plan A.1, 2026-09-14): joint stream + lean shared read-out + onset-hazard gated anchor mixture; flags via cwm3d_v2.add_args | |
| _formulation_kw(name, a, False) | |
| from aide.models import cwm3d_spatial, cwm3d_v2, heads, motionbert_kp | |
| return cwm3d_v2.build(name, **cwm3d_v2.kwargs_from_args(a)), [cwm3d_v2.__file__, heads.__file__, cwm3d_spatial.__file__, | |
| motionbert_kp.__file__, LS.__file__] | |
| if name.startswith("cwm3d"): | |
| from aide.models import cwm3d, heads | |
| kw = dict(d=a.d, layers=a.layers, heads=a.heads, ffn=a.ffn, head_hidden=a.head_hidden, dropout=a.dropout) | |
| if name == "cwm3d_kp_spatial": # hybrid: DSTformer-style spatial attention before the part tokenizer (coordinator 2026-09-12) | |
| _formulation_kw(name, a, False) | |
| from aide.models import cwm3d_spatial, motionbert_kp | |
| return cwm3d_spatial.build(name, **kw), [cwm3d_spatial.__file__, cwm3d.__file__, heads.__file__, motionbert_kp.__file__, LS.__file__] | |
| kw.update(_formulation_kw(name, a, True)) # {} unless --formulation abs | |
| if a.no_anchor: | |
| kw["anchor"] = False | |
| if a.no_uncertainty: | |
| kw["uncertainty"] = False | |
| if a.dct: | |
| kw["dct"] = True | |
| if a.rollout: | |
| kw["rollout"] = a.rollout | |
| if getattr(a, "visual", False): | |
| kw.update(use_visual=True, traffic_gate=bool(getattr(a, "traffic_gate", False)), | |
| modality_dropout=float(getattr(a, "modality_dropout", 0.3))) | |
| if getattr(a, "semantic", False): | |
| kw["use_semantic"] = True | |
| kw["sem_input"] = getattr(a, "sem_input", "tokens") | |
| kw["sem_view_gate"] = bool(getattr(a, "sem_view_gate", False)) | |
| kw["sem_frame_avg"] = tuple(x for x in getattr(a, "sem_frame_avg", "").split(",") if x) | |
| if getattr(a, "egomotion", False): | |
| kw["use_ego"] = True | |
| if getattr(a, "refine_late", False): | |
| kw["refine_late"] = True | |
| if hg not in (None, "none"): | |
| kw["horizon_gain"] = hg # pre-registered arms (b)/(c); default none -> build call unchanged | |
| return cwm3d.build(name, **kw), [cwm3d.__file__, heads.__file__, LS.__file__] | |
| if name.startswith("driverwm_kp"): | |
| # arch attribution (2), decisions 2026-09-14: Driver-WM reproduction whose internal stream reads cwm3d's keypoint part tokens | |
| _formulation_kw(name, a, False) | |
| from aide.models import driverwm_kp, driverwm_repro, heads | |
| return driverwm_kp.build("driverwm_kp", dropout=a.dropout), [driverwm_kp.__file__, driverwm_repro.__file__, heads.__file__] | |
| if name.startswith("driverwm_repro"): | |
| _formulation_kw(name, a, False) | |
| from aide.models import driverwm_repro | |
| return driverwm_repro.build("driverwm_repro", dropout=a.dropout, lat_norm=getattr(a, "repro_lat_norm", "sum"), | |
| use_semantic=True), [driverwm_repro.__file__] | |
| raise ValueError(name) | |
| def is_repro(name: str) -> bool: | |
| return name.startswith("driverwm_repro") | |
| def needs_feats(a) -> bool: | |
| return bool(getattr(a, "visual", False)) or is_repro(a.model) or a.model == "sem_only" | |
| def load_feats(a, split_masks: dict, clip_ids_cache): | |
| """feats_qwen.npz[vis_key] -> {split: tensor [n,10,4,2048] f16 (cpu)}; asserts cache alignment.""" | |
| z = np.load(a.feats, allow_pickle=False) | |
| assert np.array_equal(z["clip_id"].astype(str), clip_ids_cache.astype(str)), "feature rows not aligned to the cache" | |
| assert tuple(z["views"].tolist()) == tuple(P.VIEWS) | |
| F = z[a.vis_key] | |
| return {k: torch.as_tensor(F[m]) for k, m in split_masks.items()} | |
| def load_ego(a, split_masks: dict, clip_ids_cache): | |
| """feats_egomotion.npz -> {split: tensor [n,10,3,16] f32 (cpu)}; asserts cache alignment.""" | |
| z = np.load(a.ego_feats, allow_pickle=False) | |
| assert np.array_equal(z["clip_id"].astype(str), clip_ids_cache.astype(str)), "ego rows not aligned to the cache" | |
| E = z["ego"].astype(np.float32) | |
| return {k: torch.as_tensor(E[m]) for k, m in split_masks.items()} | |
| def nearest_slot(grids: np.ndarray) -> np.ndarray: | |
| """[B,10] frame indices -> [B,10] index of the nearest protocol frame (feature slot).""" | |
| return np.abs(grids[..., None] - PROTO_IDX[None, None, :]).argmin(-1) | |
| def model_pred(out): | |
| return out["pred"] if isinstance(out, dict) else out | |
| def rank_auc(score, y): | |
| """Mann-Whitney AUC with average ranks for ties (K-v2 val diagnostic `val_haz_auc_h3`, logged only); nan if one class is absent.""" | |
| score, y = np.asarray(score, np.float64), np.asarray(y).astype(bool) | |
| n_pos, n_neg = int(y.sum()), int((~y).sum()) | |
| if n_pos == 0 or n_neg == 0: | |
| return float("nan") | |
| order = np.argsort(score, kind="mergesort") | |
| _, inv, cnt = np.unique(score[order], return_inverse=True, return_counts=True) | |
| avg = np.bincount(inv, weights=np.arange(1, len(score) + 1, dtype=np.float64)) / cnt | |
| ranks = np.empty(len(score)); ranks[order] = avg[inv] | |
| return float((ranks[y].sum() - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg)) | |
| # ----------------------------------------------------------------------------- data | |
| def random_grid(rng: np.random.Generator, n=P.N_FRAMES_STD): | |
| """10 indices with gaps in {4,5}; span = 45 - g where g = #gaps of 4 (>=1), start in [0, g-1].""" | |
| g = int(rng.integers(1, P.N_SAMPLE)) # 1..9 gaps of 4 | |
| gaps = np.array([4] * g + [5] * (P.N_SAMPLE - 1 - g)) | |
| rng.shuffle(gaps) | |
| start = int(rng.integers(0, g)) # 0..g-1 | |
| idx = start + np.concatenate([[0], np.cumsum(gaps)]) | |
| assert idx[-1] <= n - 1 | |
| return idx | |
| def random_grids(rng: np.random.Generator, B: int, p_proto=0.5, n=P.N_FRAMES_STD): | |
| """[B,10] grids: with prob p_proto the protocol grid, else random_grid() -- fully vectorised.""" | |
| g = rng.integers(1, P.N_SAMPLE, size=B) # gaps of 4 per row | |
| ranks = np.argsort(np.argsort(rng.random((B, P.N_SAMPLE - 1)), 1), 1) # random permutation ranks | |
| gaps = np.where(ranks < g[:, None], 4, 5) | |
| start = np.floor(rng.random(B) * g).astype(int) | |
| idx = start[:, None] + np.concatenate([np.zeros((B, 1), int), np.cumsum(gaps, 1)], 1) | |
| assert idx[:, -1].max() <= n - 1 | |
| proto = rng.random(B) < p_proto | |
| idx[proto] = P.sample_idx(n) | |
| return idx | |
| def split_obs_fut(frames: torch.Tensor): | |
| """frames [B,10,J,3] -> obs [B,5,J,3], fut_xy [B,5,J,2]. The ONLY place targets are cut.""" | |
| return frames[:, :P.T_OBS], frames[:, P.T_OBS:, :, :2] | |
| class TrainSampler: | |
| def __init__(self, S45: np.ndarray, aug_window: bool, seed: int, noise_px=0.0, joint_drop=0.0, | |
| translate=0.0, scale=0.0, flip_p=0.0, device=None, feats=None, labels=None, ego=None, hazard_tau=None): | |
| dev = torch.device("cpu") if device is None else device | |
| self.S = torch.as_tensor(S45).to(dev) # [N,45,J,3] resident on the device (138 MB) | |
| # K-v2: per-part displacement thresholds [4] px (TRAIN statistics) -> onset / risk targets emitted by batch_extra; None = no targets | |
| self.hazard_tau = None if hazard_tau is None else torch.as_tensor(hazard_tau, dtype=torch.float32).to(dev) | |
| self.feats = None if feats is None else feats.to(dev) # [N,10,4,2048] f16 (protocol slots) | |
| self.ego = None if ego is None else ego.to(dev) # [N,10,3,16] f32 (protocol slots) | |
| self.labels = None if labels is None else {k: torch.as_tensor(v).to(dev) for k, v in labels.items()} | |
| self.aug_window, self.noise_px, self.joint_drop = aug_window, noise_px, joint_drop | |
| self.translate, self.scale, self.flip_p = translate, scale, flip_p | |
| self.rng = np.random.default_rng(seed) | |
| self.proto = P.sample_idx(P.N_FRAMES_STD) | |
| self.flip = torch.as_tensor(P.flip_perm()).to(dev) | |
| def batch(self, idx: np.ndarray, device): | |
| obs, fut, _ = self.batch_extra(idx, device) | |
| return obs, fut | |
| def batch_extra(self, idx: np.ndarray, device): | |
| """-> obs [B,5,J,3], fut_xy [B,5,J,2], extra {vis [B,10,4,2048] (if feats), y {task: [B]}, fut_conf [B,5,J]}.""" | |
| B = len(idx) | |
| grids = random_grids(self.rng, B) if self.aug_window else np.tile(self.proto, (B, 1)) # [B,10] | |
| ti = torch.as_tensor(idx, device=self.S.device) | |
| x = self.S[ti[:, None], torch.as_tensor(grids, device=self.S.device)] | |
| x = x.to(device, non_blocking=True).float().clone() # [B,10,J,3] | |
| extra = {} | |
| if self.feats is not None or self.ego is not None: | |
| slot = torch.as_tensor(nearest_slot(grids), device=self.S.device) | |
| if self.feats is not None: | |
| extra["vis"] = self.feats[ti[:, None], slot].to(device, non_blocking=True) # [B,10,4,2048] | |
| if self.ego is not None: | |
| extra["ego"] = self.ego[ti[:, None], slot].to(device, non_blocking=True) # [B,10,3,16] | |
| if self.labels is not None: | |
| extra["y"] = {k: v[ti].to(device) for k, v in self.labels.items()} | |
| # geometric (all 10 frames, consistent) | |
| if self.scale > 0: | |
| s = 1 + (torch.rand(B, 1, 1, 1, device=device) * 2 - 1) * self.scale | |
| c = x[:, P.T_OBS - 1, :, :2].mean(1)[:, None, None] # last-observed centroid | |
| x[..., :2] = c + (x[..., :2] - c) * s | |
| if self.translate > 0: | |
| x[..., :2] += (torch.rand(B, 1, 1, 2, device=device) * 2 - 1) * self.translate | |
| if self.flip_p > 0: | |
| f = torch.rand(B, device=device) < self.flip_p | |
| if f.any(): | |
| xf = x[f][:, :, self.flip] | |
| xf[..., 0] = P.W - xf[..., 0] | |
| x[f] = xf | |
| # observation noise / missing joints (observed frames only) | |
| if self.noise_px > 0: | |
| n = self.noise_px * (1 - x[:, :P.T_OBS, :, 2:3]) * torch.randn_like(x[:, :P.T_OBS, :, :2]) | |
| x[:, :P.T_OBS, :, :2] += n | |
| if self.joint_drop > 0: | |
| drop = torch.rand(B, P.T_OBS, P.J, device=device) < self.joint_drop | |
| last_xy = x[:, P.T_OBS - 1:P.T_OBS, :, :2].expand(-1, P.T_OBS, -1, -1) | |
| x[:, :P.T_OBS, :, :2] = torch.where(drop[..., None], last_xy, x[:, :P.T_OBS, :, :2]) | |
| x[:, :P.T_OBS, :, 2] = torch.where(drop, torch.zeros_like(x[:, :P.T_OBS, :, 2]), x[:, :P.T_OBS, :, 2]) | |
| extra["fut_conf"] = x[:, P.T_OBS:, :, 2] | |
| obs, fut = split_obs_fut(x) | |
| if self.hazard_tau is not None: # K-v2 onset-hazard targets from the SAME future block (no other cut point) | |
| extra["haz_onset"], extra["haz_risk"] = LS.hazard_targets(obs[:, -1, :, :2], fut[..., :2], self.hazard_tau.to(device)) | |
| return obs, fut, extra | |
| PRED_EXTRA_KEYS = ("haz_P",) # additional model outputs saved next to `pred` in pred_<split>.npz (the harness ignores unknown keys) | |
| def predict(model, S10: np.ndarray, device, bs=256, vis=None, repro=False, return_logits=False, intervene=None, ego=None, | |
| model_kw=None, extra_out=None): | |
| """vis: [N,10,4,2048], ego: [N,10,3,16] (only frames < T_OBS are handed to the model). intervene: callable(vis_obs)->vis_obs. | |
| model_kw: extra keyword arguments forwarded to the model call (test-time interventions: ext_scale, gate_override, | |
| repro `override`); None = the plain forward used for every reported run. | |
| extra_out: optional dict filled with the concatenated PRED_EXTRA_KEYS outputs present in the model dict (K-v2 haz_P [N,4,H]).""" | |
| model.eval() | |
| model_kw = dict(model_kw or {}) | |
| preds, logits, extras = [], {}, {} | |
| for i in range(0, len(S10), bs): | |
| obs, _ = split_obs_fut(torch.as_tensor(S10[i:i + bs]).to(device).float()) | |
| v = None | |
| if vis is not None: | |
| v = torch.as_tensor(vis[i:i + bs]).to(device)[:, :P.T_OBS] | |
| if intervene is not None: | |
| v = intervene(v) | |
| e = None if ego is None else torch.as_tensor(ego[i:i + bs]).to(device)[:, :P.T_OBS] | |
| if repro: | |
| out = model(v, **model_kw) | |
| else: | |
| kw = dict(model_kw) | |
| if v is not None: | |
| kw["vis"] = v | |
| if e is not None: | |
| kw["ego"] = e | |
| out = model(obs, **kw) | |
| preds.append(model_pred(out).float().cpu().numpy()) | |
| if isinstance(out, dict): | |
| for t in P.TASKS: | |
| if f"logits_{t}" in out: | |
| logits.setdefault(t, []).append(out[f"logits_{t}"].float().cpu().numpy()) | |
| if extra_out is not None: | |
| for k in PRED_EXTRA_KEYS: | |
| if k in out and out[k] is not None: | |
| extras.setdefault(k, []).append(out[k].float().cpu().numpy()) | |
| pred = np.concatenate(preds) | |
| if extra_out is not None: | |
| extra_out.update({k: np.concatenate(v) for k, v in extras.items()}) | |
| if return_logits: | |
| return pred, {t: np.concatenate(v) for t, v in logits.items()} | |
| return pred | |
| class EMA: | |
| def __init__(self, model, decay): | |
| self.decay = decay | |
| self.shadow = copy.deepcopy(model).eval() | |
| for p in self.shadow.parameters(): | |
| p.requires_grad_(False) | |
| def update(self, model): | |
| for ps, pm in zip(self.shadow.parameters(), model.parameters()): | |
| ps.mul_(self.decay).add_(pm.detach(), alpha=1 - self.decay) | |
| for bs_, bm in zip(self.shadow.buffers(), model.buffers()): | |
| bs_.copy_(bm) | |
| def seed_all(seed): | |
| random.seed(seed) | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| torch.cuda.manual_seed_all(seed) | |
| # ----------------------------------------------------------------------------- main | |
| def get_args(argv=None): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--model", default="simlpe_kp") | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--name", default=None) | |
| ap.add_argument("--epochs", type=int, default=200) | |
| ap.add_argument("--bs", type=int, default=32) | |
| ap.add_argument("--lr", type=float, default=3e-4) | |
| ap.add_argument("--wd", type=float, default=0.05) | |
| ap.add_argument("--warmup", type=int, default=10, help="warm-up epochs") | |
| ap.add_argument("--ema", type=float, default=0.999) | |
| ap.add_argument("--patience", type=int, default=40) | |
| ap.add_argument("--clip", type=float, default=1.0) | |
| # losses | |
| ap.add_argument("--loss", default="euc", choices=("euc", "l1", "mse"), help="plain-output models") | |
| ap.add_argument("--w-euc", type=float, default=1.0) | |
| ap.add_argument("--w-nll", type=float, default=0.1) | |
| ap.add_argument("--w-bone", type=float, default=0.05) | |
| ap.add_argument("--w-jerk", type=float, default=0.01) | |
| ap.add_argument("--no-physics", action="store_true", help="bone and jerk weights -> 0") | |
| # augmentation | |
| ap.add_argument("--aug-window", type=int, default=1) | |
| ap.add_argument("--noise-px", type=float, default=0.0) | |
| ap.add_argument("--joint-drop", type=float, default=0.0) | |
| ap.add_argument("--translate", type=float, default=0.0) | |
| ap.add_argument("--scale", type=float, default=0.0) | |
| ap.add_argument("--flip-p", type=float, default=0.0) | |
| ap.add_argument("--dropout", type=float, default=0.0) | |
| # cwm3d architecture flags | |
| ap.add_argument("--d", type=int, default=256) | |
| ap.add_argument("--layers", type=int, default=4) | |
| ap.add_argument("--heads", type=int, default=4) | |
| ap.add_argument("--ffn", type=int, default=1024) | |
| ap.add_argument("--head-hidden", type=int, default=512) | |
| ap.add_argument("--no-anchor", action="store_true") | |
| ap.add_argument("--no-uncertainty", action="store_true") | |
| ap.add_argument("--dct", action="store_true") | |
| ap.add_argument("--rollout", default=None, choices=(None, "chunked", "recursive")) | |
| # visual tokens / semantic heads (Stage B) and the Driver-WM reproduction (Stage C) | |
| ap.add_argument("--visual", action="store_true", help="cwm3d: Qwen3-VL tokens per observed frame x view") | |
| ap.add_argument("--vis-key", default="pooled", choices=("pooled", "vis_pooled")) | |
| ap.add_argument("--feats", default=str(FEATS_NPZ)) | |
| ap.add_argument("--traffic-gate", action="store_true", help="cwm3d: traffic->driver gated cross-attention") | |
| ap.add_argument("--modality-dropout", type=float, default=0.3) | |
| ap.add_argument("--semantic", action="store_true", help="cwm3d: 4 semantic heads (CE)") | |
| ap.add_argument("--lambda-sem", type=float, default=0.5, help="weight per semantic head") | |
| ap.add_argument("--label-smoothing", type=float, default=0.1) | |
| ap.add_argument("--class-balanced", type=int, default=1, help="inverse-frequency class weights (train)") | |
| ap.add_argument("--egomotion", action="store_true", help="cwm3d: ego-motion (optical-flow) token per observed frame") | |
| ap.add_argument("--ego-feats", default=str(EGO_NPZ)) | |
| ap.add_argument("--sem-input", default="tokens", choices=("tokens", "direct", "attn"), | |
| help="semantic heads read 256-d tokens (plan), the frozen 2048-d features directly, or a 1-layer " | |
| "causal attention pooling over the window features (repro recipe)") | |
| ap.add_argument("--select", default="mpjpe", choices=("mpjpe", "f1"), | |
| help="checkpoint selection: val All-MPJPE (default) or val mean macro-F1 (semantic-only models)") | |
| ap.add_argument("--sem-task-weights", default="", help="per-head CE multipliers, e.g. der=2") | |
| ap.add_argument("--sem-view-gate", action="store_true", help="attn heads: learned softmax gate over the 4 view encodings") | |
| ap.add_argument("--sem-frame-avg", default="", help="attn heads: tasks whose logits are averaged over observed frames, e.g. der") | |
| ap.add_argument("--dual-select", action="store_true", | |
| help="also keep the checkpoint with the best val mean macro-F1 and evaluate it as <name>__f1sel") | |
| ap.add_argument("--h-weights", default="", help="horizon weights for euc/NLL, e.g. 1,1,1,2,3 (normalised to mean 1)") | |
| ap.add_argument("--motion-weight", type=float, default=0.0, | |
| help="per-clip weight = min(GT future displacement / batch mean, cap); 0 = off") | |
| ap.add_argument("--motion-oversample", type=float, default=0.0, | |
| help="roadmap B1: per-epoch sampling of TRAIN clips with replacement, p_i ∝ 1 + k·s_i/mean(s), " | |
| "s_i = GT future displacement of the clip on the protocol frames (HM motion score); 0 = uniform permutation") | |
| ap.add_argument("--patience-f1", type=int, default=0, | |
| help="roadmap A1: with --dual-select, early-stop only when BOTH val All-MPJPE (patience) and val mean " | |
| "macro-F1 (this many epochs) have stalled; 0 = geometry patience only (old behaviour)") | |
| ap.add_argument("--refine-late", action="store_true", | |
| help="cwm3d: zero-initialised late-horizon refinement head (Δ for h = 4, 5 from last-frame part tokens " | |
| "+ detached h <= 3 predictions; roadmap B3)") | |
| ap.add_argument("--repro-lat-norm", default="sum", choices=("sum", "mean")) | |
| ap.add_argument("--init-ckpt", default=None, | |
| help="opt-in (GPU-baselines A, protocol C21): initialise the model from this checkpoint's `model` state dict before training " | |
| "(e.g. aide/results/runs/cwm3d_kp_batonpre_pretrain_s42/ckpt_best.pt); tensors are matched by name AND shape, the rest keep " | |
| "the seed init; default None = unchanged behaviour (seed init, exact zero-velocity start)") | |
| ap.add_argument("--init-parts", default="all", choices=("all", "trunk"), | |
| help="with --init-ckpt: `all` = every matching tensor (trunk + horizon head), `trunk` = tokenizer / blocks / ln_f only " | |
| "(the horizon head keeps its zero init, so the fine-tune still starts at zero-velocity)") | |
| ap.add_argument("--formulation", default="residual", choices=("residual", "abs"), | |
| help="docs/decisions.md 2026-09-15 formulation ablation (explanatory, not for adoption): `residual` = unchanged behaviour " | |
| "(pred = last pose + 200·out, zero-init read-out; cwm3d: anchor mixture); `abs` = the read-out emits ABSOLUTE normalised " | |
| "coordinates, pred = out ⊙ (W, H) px, PyTorch default (non-zero) init, no residual to the last pose, no anchors — " | |
| "motionbert_kp* and cwm3d* (aide/models/cwm3d.py) only; every other model rejects it") | |
| # pre-registered HM h=5 arms (docs/decisions.md 2026-09-17T07:00:22Z, part (A)); defaults = unchanged behaviour | |
| ap.add_argument("--hand-weight", type=float, default=1.0, | |
| help="arms (a1)/(a2): weight of the 42 hand keypoints (PARTS lhand + rhand, HALPE 94-135) in the Euclidean loss term, other " | |
| "joints 1, not renormalised; NLL / bone / jerk terms and every metric unchanged; 1.0 = the untouched code path") | |
| ap.add_argument("--horizon-gain", default="none", choices=("none", "learned", "parts"), | |
| help="cwm3d* (aide/models/cwm3d.py) only: learned per-horizon gain g_h = 1 + δ_h (δ zero-init) multiplying the anchored residual " | |
| "(replaces the fixed SCALE): `learned` = 5 scalars (arm b), `parts` = 5 x 2 scalars other / hand joints (arm c); none = unchanged") | |
| ap.add_argument("--horizon-gain-parts", action="store_true", help="alias of --horizon-gain parts (arm c)") | |
| # K-v2 (aide/models/cwm3d_v2.py): --v2-* / --no-v2-spatial / --no-lean-head / --no-hazard / --haz-*; inert for every other model | |
| from aide.models import cwm3d_v2 | |
| cwm3d_v2.add_args(ap) | |
| # F7 ensemble distillation (pre-registered 2026-09-17; defaults = unchanged behaviour) | |
| ap.add_argument("--distil-teachers", default="", | |
| help="comma-separated run names (results/runs/<run>/ckpt_best.pt, same architecture as --model) whose MEAN prediction is the " | |
| "distillation target; '' (default) = no teacher, unchanged code path") | |
| ap.add_argument("--w-distil", type=float, default=0.0, help="weight of the coordinate distillation term (Euclidean to the mean teacher prediction); 0 = off") | |
| ap.add_argument("--w-distil-modes", type=float, default=0.0, help="cwm3d only: weight of KL(mean teacher anchor_w || student anchor_w); 0 = off") | |
| ap.add_argument("--distil-root", default=None, help="root of the teacher run dirs (default results/runs; smoke tests redirect)") | |
| ap.add_argument("--w-distil-logb", type=float, default=0.0, help="cwm3d only: weight of the L1 term between student log_b and the mean teacher log_b; 0 = off") | |
| # run control | |
| ap.add_argument("--max-minutes", type=float, default=10.0) | |
| ap.add_argument("--n-boot", type=int, default=10000) | |
| ap.add_argument("--refs", default=",".join(DEFAULT_REFS), help="name=run,... for paired Δ") | |
| ap.add_argument("--no-lock", action="store_true") | |
| ap.add_argument("--cpu", action="store_true") | |
| ap.add_argument("--notes", default="") | |
| ap.add_argument("--limit-train", type=int, default=0, help="DEBUG: use only the first N train clips (shape tests)") | |
| ap.add_argument("--out-root", default=None, help="run dir root (default results/runs)") | |
| ap.add_argument("--no-csv", action="store_true", help="do not append harness rows to results.csv") | |
| ap.add_argument("--defer-eval", action="store_true", help="train+predict only; write pending_eval.json") | |
| ap.add_argument("--eval-only", action="store_true", help="only run the deferred evaluation of --name") | |
| a = ap.parse_args(argv) | |
| if a.horizon_gain_parts: | |
| a.horizon_gain = "parts" | |
| if a.model.startswith("driverwm_kp"): | |
| a.visual = True # the exterior-view Qwen features are loaded and handed to the model exactly as for cwm3d --visual | |
| if a.model.startswith("cwm3d") and a.model != "cwm3d_plain": | |
| # plan §4 recipe defaults unless overridden explicitly | |
| defaults = {"noise_px": 2.0, "joint_drop": 0.1, "translate": 20.0, "scale": 0.05} | |
| for k, v in defaults.items(): | |
| if getattr(a, k) == 0.0 and f"--{k.replace('_', '-')}" not in (argv or sys.argv): | |
| setattr(a, k, v) | |
| if a.no_physics: | |
| a.w_bone, a.w_jerk = 0.0, 0.0 | |
| if a.no_uncertainty: | |
| a.w_nll = 0.0 | |
| return a | |
| def load_init_ckpt(model, path, parts="all", device="cpu"): | |
| """--init-ckpt: load the matching tensors of a checkpoint's `model` state dict (name AND shape must match; `trunk` = tok. / blocks. / ln_f. | |
| for cwm3d, or the model's own `TRUNK_PREFIXES` (motionbert_kp / stgcn_kp: input projection, embeddings, blocks[, ln_f]; read-out head kept)). Returns a report for the run json (what was loaded / skipped); never raises on extra or missing keys (strict=False).""" | |
| ck = torch.load(path, map_location=device, weights_only=False) | |
| sd = ck["model"] if isinstance(ck, dict) and "model" in ck else ck | |
| own = model.state_dict() | |
| trunk_prefixes = tuple(getattr(model, "TRUNK_PREFIXES", ("tok.", "blocks.", "ln_f."))) # cwm3d default; baselines declare their own | |
| load, skipped_shape, skipped_part = {}, [], [] | |
| for k, v in sd.items(): | |
| if k not in own: | |
| continue | |
| if parts == "trunk" and not k.startswith(trunk_prefixes): | |
| skipped_part.append(k) | |
| continue | |
| if tuple(own[k].shape) != tuple(v.shape): | |
| skipped_shape.append(k) | |
| continue | |
| load[k] = v | |
| missing = [k for k in own if k not in load] | |
| model.load_state_dict(load, strict=False) | |
| import hashlib | |
| h = hashlib.sha256() | |
| for k in sorted(load): | |
| h.update(k.encode()); h.update(load[k].detach().cpu().contiguous().numpy().tobytes()) | |
| return {"ckpt": str(path), "parts": parts, "ckpt_epoch": ck.get("epoch") if isinstance(ck, dict) else None, "n_loaded": len(load), "n_skipped_shape": len(skipped_shape), | |
| "n_skipped_part": len(skipped_part), "n_missing_or_kept": len(missing), "skipped_shape": skipped_shape[:20], "kept_keys": missing[:40], | |
| "loaded_sha256": h.hexdigest()[:16], "ckpt_sha256": __import__("common.runmeta", fromlist=["sha256_file"]).sha256_file(path)[:16]} | |
| def load_teachers(a, device, runs_root=None): | |
| """--distil-teachers: build `a.model` once per teacher run (same architecture flags as the student), load results/runs/<run>/ckpt_best.pt | |
| `model` state dict STRICT (EMA weights), eval mode, no gradients. Returns (teachers, report) or (None, None) when no teacher / all weights 0. | |
| Never touches anything but the teachers' ckpt_best.pt files.""" | |
| names = [x for x in (getattr(a, "distil_teachers", "") or "").split(",") if x] | |
| w_on = any(float(getattr(a, k, 0.0)) > 0 for k in ("w_distil", "w_distil_modes", "w_distil_logb")) | |
| if not names or not w_on: | |
| assert not (names and not w_on), "--distil-teachers given but every --w-distil* weight is 0" | |
| assert not (w_on and not names), "--w-distil* > 0 needs --distil-teachers" | |
| return None, None | |
| assert not getattr(a, "visual", False) and not getattr(a, "semantic", False) and not getattr(a, "egomotion", False), \ | |
| "F7 distillation is implemented for keypoint-only students (no vis / semantic / ego inputs)" | |
| root = Path(runs_root) if runs_root else P.RESULTS / "runs" | |
| from common.runmeta import sha256_file | |
| teachers, rep = [], [] | |
| for n in names: | |
| ck_path = root / n / "ckpt_best.pt" | |
| ck = torch.load(ck_path, map_location=device, weights_only=False) | |
| t, _ = build_model(a.model, a) | |
| t.load_state_dict(ck["model"], strict=True) | |
| t = t.to(device).eval() | |
| for p_ in t.parameters(): | |
| p_.requires_grad_(False) | |
| teachers.append(t) | |
| rep.append({"run": n, "ckpt": str(ck_path), "ckpt_epoch": ck.get("epoch"), "ckpt_sha256": sha256_file(ck_path)[:16]}) | |
| report = {"teachers": rep, "n_teachers": len(teachers), "w_distil": float(a.w_distil), "w_distil_modes": float(a.w_distil_modes), | |
| "w_distil_logb": float(a.w_distil_logb), "target": "mean over teachers of pred (and log_b / anchor_w when the outputs are dicts), " | |
| "computed on the SAME augmented observed window as the student, no grad"} | |
| return teachers, report | |
| def teacher_target(teachers, obs): | |
| """mean teacher outputs on obs -> {"pred": [B,H,J,2], "log_b": [B,H,J] or None, "anchor_w": [B,H,J,3] or None}.""" | |
| outs = [t(obs) for t in teachers] | |
| preds = torch.stack([model_pred(o) for o in outs]).mean(0) | |
| tgt = {"pred": preds, "log_b": None, "anchor_w": None} | |
| if all(isinstance(o, dict) for o in outs): | |
| if all(o.get("log_b") is not None for o in outs): | |
| tgt["log_b"] = torch.stack([o["log_b"] for o in outs]).mean(0) | |
| if all(o.get("anchor_w") is not None for o in outs): | |
| tgt["anchor_w"] = torch.stack([o["anchor_w"] for o in outs]).mean(0) | |
| return tgt | |
| def run_dir_of(a, name): | |
| return (Path(a.out_root) if a.out_root else P.RESULTS / "runs") / name | |
| def main(argv=None): | |
| a = get_args(argv) | |
| name = a.name or f"{a.model}_s{a.seed}" | |
| out_dir = run_dir_of(a, name) | |
| if a.eval_only: | |
| return finish_eval(out_dir) | |
| info = train_and_predict(a, name, out_dir) | |
| if a.defer_eval: | |
| (out_dir / "pending_eval.json").write_text(json.dumps(info, indent=1, default=str)) | |
| print(f"[{name}] evaluation deferred -> {out_dir / 'pending_eval.json'}", flush=True) | |
| return info | |
| return evaluate_and_record(info) | |
| def train_and_predict(a, name, out_dir): | |
| out_dir.mkdir(parents=True, exist_ok=True) | |
| t_start = time.time() | |
| seed_all(a.seed) | |
| c = np.load(P.CACHE_NPZ, allow_pickle=False) | |
| tr, va, te = (c["split"] == s for s in ("train", "val", "test")) | |
| S45_tr = c["S"][tr] | |
| if a.limit_train: | |
| S45_tr = S45_tr[:a.limit_train] | |
| S10_va, S10_te = c["S10"][va], c["S10"][te] | |
| id_va, id_te = c["clip_id"][va], c["clip_id"][te] | |
| repro = is_repro(a.model) | |
| sem_only = a.model == "sem_only" | |
| if sem_only: | |
| a.semantic, a.visual = True, True | |
| use_sem = bool(a.semantic) or repro | |
| task_w = {t: 1.0 for t in P.TASKS} | |
| for kv in [x for x in a.sem_task_weights.split(",") if x]: | |
| k, v_ = kv.split("=") | |
| task_w[k] = float(v_) | |
| feats = load_feats(a, {"train": tr, "val": va, "test": te}, c["clip_id"]) if needs_feats(a) else None | |
| if feats is not None and a.limit_train: | |
| feats["train"] = feats["train"][:a.limit_train] | |
| egof = load_ego(a, {"train": tr, "val": va, "test": te}, c["clip_id"]) if getattr(a, "egomotion", False) else None | |
| if egof is not None and a.limit_train: | |
| egof["train"] = egof["train"][:a.limit_train] | |
| labels_tr = {t: c[f"y_{t}"][tr][:len(S45_tr)] for t in P.TASKS} if use_sem else None | |
| y_va = {t: c[f"y_{t}"][va] for t in P.TASKS} | |
| class_w = None | |
| if use_sem and a.class_balanced and not repro: | |
| class_w = {} | |
| for t in P.TASKS: | |
| cnt = np.bincount(labels_tr[t], minlength=len(P.CLASSES[t])).astype(np.float64) | |
| wgt = len(labels_tr[t]) / (len(P.CLASSES[t]) * np.maximum(cnt, 1)) | |
| class_w[t] = wgt / wgt.mean() | |
| fut_va = S10_va[:, P.T_OBS:] | |
| zv_va = M.mpjpe(M.joint_errors(np.repeat(S10_va[:, P.T_OBS - 1:P.T_OBS, :, :2], P.T_PRED, 1), fut_va)) | |
| hm_va, _, _ = M.hm_mask(fut_va) | |
| lock = None | |
| if not a.no_lock and not a.cpu: | |
| lock = open(GPU_LOCK, "a") | |
| t_l = time.time() | |
| fcntl.flock(lock, fcntl.LOCK_EX) | |
| print(f"GPU lock acquired after {time.time() - t_l:.1f}s", flush=True) | |
| t_gpu = time.time() | |
| device = torch.device("cpu" if a.cpu else "cuda") | |
| model, code_files = build_model(a.model, a) | |
| if egof is not None: | |
| flat = egof["train"].flatten(0, 1).flatten(1) # [n*10, 48] train statistics for standardisation | |
| model.set_ego_stats(flat.mean(0), flat.std(0)) | |
| model = model.to(device) | |
| init_report = None | |
| if getattr(a, "init_ckpt", None): # opt-in pretrained init (default None: unchanged) | |
| init_report = load_init_ckpt(model, a.init_ckpt, getattr(a, "init_parts", "all"), device) | |
| print(f"[{name}] init from {a.init_ckpt} ({a.init_parts}): {init_report['n_loaded']} tensors loaded, " | |
| f"{init_report['n_skipped_shape']} shape-mismatched skipped, {init_report['n_missing_or_kept']} kept at seed init", flush=True) | |
| teachers, distil_report = load_teachers(a, device, getattr(a, "distil_root", None)) # None unless --distil-teachers (F7) | |
| if teachers is not None: | |
| print(f"[{name}] F7 distillation: {len(teachers)} teachers {[t['run'] for t in distil_report['teachers']]} " | |
| f"w coord {a.w_distil} modes {a.w_distil_modes} log_b {a.w_distil_logb}", flush=True) | |
| hazard_info, haz_tau, haz_val = None, None, None | |
| if getattr(model, "haz", None) is not None: | |
| # K-v2: displacement thresholds tau_p and the cumulative onset prior from the TRAIN split only (protocol frames of the 45-frame | |
| # windows, as --motion-oversample); the prior initialises the hazard bias, so the untrained head outputs P == prior for every clip | |
| tau, prior = LS.hazard_stats(S45_tr[:, PROTO_IDX], a.haz_quantile) | |
| model.set_hazard_prior(tau, prior) | |
| haz_tau = tau | |
| hazard_info = {"tau": tau.tolist(), "prior_P": prior.tolist(), "quantile": float(a.haz_quantile), "n_train_clips": int(len(S45_tr)), | |
| "weight": float(a.haz_weight), "modulation": model.mod is not None} | |
| x_va = torch.as_tensor(S10_va).float() # val exceedance labels for the logged AUC diagnostic only | |
| haz_val = torch.cummax(LS.hazard_targets(x_va[:, P.T_OBS - 1, :, :2], x_va[:, P.T_OBS:, :, :2], tau)[0], -1).values.numpy() | |
| print(f"[{name}] hazard prior (train, q={a.haz_quantile}): tau {np.round(tau, 1).tolist()} px, P_h5 {np.round(prior[:, -1], 3).tolist()}", flush=True) | |
| n_params = sum(p.numel() for p in model.parameters()) | |
| ema = EMA(model, a.ema) | |
| opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.wd) | |
| steps_per_epoch = math.ceil(len(S45_tr) / a.bs) | |
| total = a.epochs * steps_per_epoch | |
| warm = max(1, a.warmup * steps_per_epoch) | |
| sched = torch.optim.lr_scheduler.LambdaLR( | |
| opt, lambda s: (s + 1) / warm if s < warm else 0.5 * (1 + math.cos(math.pi * (s - warm) / max(1, total - warm)))) | |
| sampler = TrainSampler(S45_tr, bool(a.aug_window), a.seed, a.noise_px, a.joint_drop, a.translate, a.scale, a.flip_p, | |
| device=device, feats=None if feats is None else feats["train"], labels=labels_tr, | |
| ego=None if egof is None else egof["train"], hazard_tau=haz_tau) | |
| ego_va = None if egof is None else egof["val"].to(device) | |
| ego_te = None if egof is None else egof["test"].to(device) | |
| class_wt = None if class_w is None else {t: torch.as_tensor(w, dtype=torch.float32, device=device) for t, w in class_w.items()} | |
| vis_va = None if feats is None else feats["val"].to(device) | |
| vis_te = None if feats is None else feats["test"].to(device) | |
| rng = np.random.default_rng(a.seed) | |
| w = {"euc": a.w_euc, "nll": a.w_nll, "bone": a.w_bone, "jerk": a.w_jerk} | |
| h_weights = [float(x) for x in a.h_weights.split(",")] if a.h_weights else None | |
| jw = LS.joint_weights(getattr(a, "hand_weight", 1.0), device=device) # None unless --hand-weight != 1 (arms a1/a2/d) | |
| best_f1, best_f1_ep, bad_f1 = -1.0, -1, 0 | |
| sample_p, oversample_report = None, None | |
| if a.motion_oversample and a.motion_oversample > 0: | |
| # per-clip motion score on the TRAIN clip's protocol frames (frames 5..9 of sample_idx(45)); labels/targets of the | |
| # train split only -- nothing from val/test enters the sampler | |
| s_tr = M.motion_score(S45_tr[:, PROTO_IDX][:, P.T_OBS:]) | |
| wts = 1.0 + a.motion_oversample * s_tr / max(float(s_tr.mean()), 1e-6) | |
| sample_p = wts / wts.sum() | |
| oversample_report = {"k": a.motion_oversample, "formula": "p_i ∝ 1 + k·s_i/mean(s)", "w_min": float(wts.min()), | |
| "w_max": float(wts.max()), "w_mean": float(wts.mean()), | |
| "expected_share_top10pct": float(np.sort(sample_p)[-max(1, len(sample_p) // 10):].sum())} | |
| hist, best, best_ep, bad, stop_reason = [], float("inf"), -1, 0, "epochs" | |
| print(f"{name}: params {n_params:,} train {len(S45_tr)} val {len(S10_va)} test {len(S10_te)} " | |
| f"val ZV All-MPJPE {zv_va:.3f} loss w {w} aug noise {a.noise_px} drop {a.joint_drop} " | |
| f"tr {a.translate} sc {a.scale}" | |
| + (f" hand_weight {a.hand_weight} horizon_gain {a.horizon_gain}" if (jw is not None or getattr(model, "hgain", None) is not None) else ""), flush=True) | |
| for ep in range(a.epochs): | |
| model.train() | |
| perm = rng.permutation(len(S45_tr)) if sample_p is None else rng.choice(len(S45_tr), size=len(S45_tr), replace=True, p=sample_p) | |
| acc, nb = {}, 0 | |
| for i in range(0, len(perm), a.bs): | |
| obs, fut, ex = sampler.batch_extra(perm[i:i + a.bs], device) | |
| if repro: | |
| out = model(ex["vis"][:, :P.T_OBS]) | |
| loss, terms = model.loss(out, fut, ex["fut_conf"], ex["vis"][:, P.T_OBS:], ex.get("y")) | |
| else: | |
| mkw = {} | |
| if a.visual: | |
| mkw["vis"] = ex["vis"][:, :P.T_OBS] | |
| if egof is not None: | |
| mkw["ego"] = ex["ego"][:, :P.T_OBS] | |
| out = model(obs, **mkw) | |
| if sem_only: | |
| loss, terms = fut.new_zeros(()), {} | |
| else: | |
| sw = LS.sample_weights(fut, h_weights, a.motion_weight) if (h_weights or a.motion_weight > 0) else None | |
| loss, terms = LS.compute_losses(out, fut, w, a.loss, sw=sw, jw=jw) | |
| if getattr(a, "haz_weight", 0.0) > 0 and isinstance(out, dict) and "haz_logits" in out and "haz_onset" in ex: # K-v2 onset-hazard BCE | |
| l_haz = LS.hazard_bce(out["haz_logits"], ex["haz_onset"], ex["haz_risk"]) | |
| loss = loss + a.haz_weight * l_haz | |
| terms = {**terms, "haz": l_haz.detach()} | |
| if teachers is not None: # F7: ensemble distillation terms (additive) | |
| l_d, d_terms = LS.distillation(out, teacher_target(teachers, obs), a.w_distil, a.w_distil_modes, a.w_distil_logb) | |
| loss = loss + l_d | |
| terms = {**terms, **d_terms} | |
| if a.semantic: | |
| l_sem = sum(task_w[t] * TF.cross_entropy(out[f"logits_{t}"], ex["y"][t], | |
| weight=None if class_wt is None else class_wt[t], | |
| label_smoothing=a.label_smoothing) for t in P.TASKS) | |
| loss = loss + a.lambda_sem * l_sem | |
| terms = {**terms, "sem": l_sem.detach()} | |
| opt.zero_grad(set_to_none=True) | |
| loss.backward() | |
| if a.clip > 0: | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), a.clip) | |
| opt.step() | |
| sched.step() | |
| ema.update(model) | |
| acc["total"] = acc.get("total", 0.0) + loss.item() | |
| for k, v in terms.items(): | |
| acc[k] = acc.get(k, 0.0) + float(v) | |
| nb += 1 | |
| ex_va = {} | |
| pv, lg = predict(ema.shadow, S10_va, device, vis=vis_va, repro=repro, return_logits=True, ego=ego_va, extra_out=ex_va) | |
| err = M.joint_errors(pv, fut_va) | |
| v_all, v_hm = M.mpjpe(err), M.mpjpe(err[hm_va]) | |
| rec = {"epoch": ep, "val_mpjpe_all": v_all, "val_mpjpe_hm": v_hm, "lr": sched.get_last_lr()[0], | |
| "t": round(time.time() - t_start, 1), **{f"train_{k}": v / nb for k, v in acc.items()}} | |
| if haz_val is not None and "haz_P" in ex_va: # logged only; selection stays val All-MPJPE | |
| rec["val_haz_auc_h3"] = float(np.nanmean([rank_auc(ex_va["haz_P"][:, p, 2], haz_val[:, p, 2]) for p in range(haz_val.shape[1])])) | |
| for t, l in lg.items(): | |
| rec[f"val_f1_{t}"] = M.macro_f1(y_va[t], l.argmax(-1)) | |
| hist.append(rec) | |
| if a.dual_select and lg: | |
| f1m = float(np.mean([rec[f"val_f1_{t}"] for t in P.TASKS if f"val_f1_{t}" in rec])) | |
| rec["val_f1_mean"] = f1m | |
| if f1m > best_f1 + 1e-6: | |
| best_f1, best_f1_ep, bad_f1 = f1m, ep, 0 | |
| torch.save({"model": ema.shadow.state_dict(), "epoch": ep, "val_f1_mean": f1m, "args": vars(a)}, | |
| out_dir / "ckpt_best_f1.pt") | |
| else: | |
| bad_f1 += 1 | |
| sel = v_all if a.select == "mpjpe" else -float(np.mean([rec[f"val_f1_{t}"] for t in P.TASKS if f"val_f1_{t}" in rec])) | |
| improved = sel < best - 1e-4 | |
| if improved: | |
| best, best_ep, bad = sel, ep, 0 | |
| torch.save({"model": ema.shadow.state_dict(), "epoch": ep, "val_mpjpe_all": v_all, "args": vars(a)}, | |
| out_dir / "ckpt_best.pt") | |
| else: | |
| bad += 1 | |
| if ep % 10 == 0 or improved or ep == a.epochs - 1: | |
| tstr = " ".join(f"{k} {v / nb:.3f}" for k, v in acc.items() if k != "total") | |
| fstr = (" F1 " + "/".join(f"{rec[f'val_f1_{t}']:.1f}" for t in P.TASKS)) if lg else "" | |
| fstr += f" hazAUC3 {rec['val_haz_auc_h3']:.3f}" if "val_haz_auc_h3" in rec else "" | |
| print(f"ep {ep:3d} loss {acc['total'] / nb:7.3f} [{tstr}] val All {v_all:7.3f} HM {v_hm:7.2f}{fstr} " | |
| f"best {best:7.3f}@{best_ep} lr {sched.get_last_lr()[0]:.2e} {(time.time() - t_start) / 60:.1f}min", | |
| flush=True) | |
| f1_patience_on = bool(a.dual_select and lg and a.patience_f1 > 0) | |
| if bad >= a.patience and (not f1_patience_on or bad_f1 >= a.patience_f1): | |
| stop_reason = (f"early_stop(patience={a.patience},patience_f1={a.patience_f1})" if f1_patience_on | |
| else f"early_stop(patience={a.patience})") | |
| break | |
| if (time.time() - t_gpu) / 60 > a.max_minutes: | |
| stop_reason = f"time_budget({a.max_minutes}min)" | |
| break | |
| ck = torch.load(out_dir / "ckpt_best.pt", map_location=device) | |
| ema.shadow.load_state_dict(ck["model"]) | |
| preds = {} | |
| for split, S10, ids, vs, eg in (("val", S10_va, id_va, vis_va, ego_va), ("test", S10_te, id_te, vis_te, ego_te)): | |
| ex_out = {} | |
| pr, lg = predict(ema.shadow, S10, device, vis=vs, repro=repro, return_logits=True, ego=eg, extra_out=ex_out) | |
| assert np.isfinite(pr).all() | |
| np.savez(out_dir / f"pred_{split}.npz", pred=pr.astype(np.float32), clip_id=ids.astype(str), split=np.array(split), | |
| **{f"logits_{t}": l.astype(np.float32) for t, l in lg.items()}, | |
| **{k: v.astype(np.float32) for k, v in ex_out.items()}) # K-v2: haz_P [N,4,H] (harness ignores extra keys) | |
| preds[split] = out_dir / f"pred_{split}.npz" | |
| f1sel = None | |
| if a.dual_select and (out_dir / "ckpt_best_f1.pt").exists(): | |
| ckf = torch.load(out_dir / "ckpt_best_f1.pt", map_location=device) | |
| ema.shadow.load_state_dict(ckf["model"]) | |
| fdir = out_dir.parent / f"{name}__f1sel" | |
| fdir.mkdir(parents=True, exist_ok=True) | |
| f1sel = {"name": f"{name}__f1sel", "out_dir": str(fdir), "epoch": best_f1_ep, "val_f1_mean": best_f1, "preds": {}} | |
| for split, S10, ids, vs, eg in (("val", S10_va, id_va, vis_va, ego_va), ("test", S10_te, id_te, vis_te, ego_te)): | |
| pr, lg = predict(ema.shadow, S10, device, vis=vs, repro=repro, return_logits=True, ego=eg) | |
| assert np.isfinite(pr).all() | |
| np.savez(fdir / f"pred_{split}.npz", pred=pr.astype(np.float32), clip_id=ids.astype(str), split=np.array(split), | |
| **{f"logits_{t}": l.astype(np.float32) for t, l in lg.items()}) | |
| f1sel["preds"][split] = str(fdir / f"pred_{split}.npz") | |
| ema.shadow.load_state_dict(ck["model"]) # back to the geometry-selected weights | |
| print(f"[{name}] dual selection: best val mean F1 {best_f1:.2f} at epoch {best_f1_ep} -> {fdir}", flush=True) | |
| gpu_minutes = (time.time() - t_gpu) / 60 | |
| if lock is not None: | |
| fcntl.flock(lock, fcntl.LOCK_UN) | |
| lock.close() | |
| (out_dir / "history.json").write_text(json.dumps(hist, indent=1)) | |
| # quick (bootstrap-free) val/test numbers so a batch driver can apply gates before the harness runs | |
| quick = {} | |
| for split, S10 in (("val", S10_va), ("test", S10_te)): | |
| pr = np.load(preds[split], allow_pickle=False)["pred"].astype(np.float64) | |
| err = M.joint_errors(pr, S10[:, P.T_OBS:]) | |
| hm, _, _ = M.hm_mask(S10[:, P.T_OBS:]) | |
| quick[split] = {"mpjpe_all": M.mpjpe(err), "mpjpe_hm": M.mpjpe(err[hm])} | |
| print(f"[{name}] quick val All {quick['val']['mpjpe_all']:.3f} HM {quick['val']['mpjpe_hm']:.2f} | " | |
| f"test All {quick['test']['mpjpe_all']:.3f} HM {quick['test']['mpjpe_hm']:.2f}; " | |
| f"GPU {gpu_minutes:.2f} min; {stop_reason}", flush=True) | |
| return {"name": name, "out_dir": str(out_dir), "args": vars(a), "code_files": [str(f) for f in code_files], | |
| "n_params": n_params, "best_epoch": best_ep, "best_val_mpjpe_all": best, "val_zero_velocity_mpjpe_all": zv_va, | |
| "epochs_run": len(hist), "stop_reason": stop_reason, "gpu_minutes": round(gpu_minutes, 2), | |
| "device": str(device), "t_start": t_start, "quick": quick, "selection": a.select, | |
| "model_report": (ema.shadow.sem.view_weights() if hasattr(ema.shadow, "sem") and hasattr(ema.shadow.sem, "view_weights") else None), | |
| "motion_oversample": oversample_report, "best_val_f1_mean": best_f1 if a.dual_select else None, "best_f1_epoch": best_f1_ep if a.dual_select else None, | |
| "preds": {k: str(v) for k, v in preds.items()}, "f1sel": f1sel, "init": init_report, "hazard": hazard_info, "distil": distil_report, | |
| "hand_weight": float(getattr(a, "hand_weight", 1.0)), | |
| "hgain": (ema.shadow.hgain.report() if getattr(ema.shadow, "hgain", None) is not None else None)} # learned gains (EMA weights) | |
| def evaluate_and_record(info: dict): | |
| """Harness (val + test, paired Δ vs ZV / refs) + run_<name>_<utc>.json. CPU only.""" | |
| a = argparse.Namespace(**info["args"]) | |
| name, out_dir = info["name"], Path(info["out_dir"]) | |
| notes = f"{a.notes};best_ep={info['best_epoch']};stop={info['stop_reason']}".strip(";") | |
| refs = parse_refs([r for r in a.refs.split(",") if r]) | |
| recs = {} | |
| for split in ("val", "test"): | |
| recs[split] = run_harness(info["preds"][split], name, a.model, a.seed, split, notes, a.n_boot, refs=refs, | |
| write_csv=not a.no_csv) | |
| print(summary_line(recs[split]), flush=True) | |
| tm, vm = recs["test"]["metrics"], recs["val"]["metrics"] | |
| keys = ("mpjpe_all", "dn_mpjpe", "mpjpe_hm", "pck05", "pck10", "mpjpe_h5", "hm_h5", "diff_vs_ref_all", | |
| "diff_vs_ref_hm", "diff_vs_zvshrink_all", "diff_vs_zvshrink_hm", "diff_vs_simlpe_all", "diff_vs_simlpe_hm", | |
| "diff_vs_kp_all", "diff_vs_kp_hm", "f1_dbr", "f1_der", "f1_tcr", "f1_vcr", "acc_dbr", "acc_der", "acc_tcr", "acc_vcr") | |
| write_run_json(out_dir, name, [__file__, P.__file__, M.__file__] + info["code_files"], [P.CACHE_NPZ], | |
| seed=a.seed, t_start=info["t_start"], input_listing=False, extra={ | |
| "config": vars(a), "n_params": info["n_params"], "best_epoch": info["best_epoch"], | |
| "best_val_mpjpe_all": info["best_val_mpjpe_all"], "selection": info.get("selection", "mpjpe"), | |
| "model_report": info.get("model_report"), "motion_oversample": info.get("motion_oversample"), | |
| "best_val_f1_mean": info.get("best_val_f1_mean"), "best_f1_epoch": info.get("best_f1_epoch"), | |
| "init": info.get("init"), "hazard": info.get("hazard"), "distil": info.get("distil"), | |
| "hand_weight": info.get("hand_weight", 1.0), "hgain": info.get("hgain"), | |
| "val_zero_velocity_mpjpe_all": info["val_zero_velocity_mpjpe_all"], | |
| "epochs_run": info["epochs_run"], "stop_reason": info["stop_reason"], | |
| "gpu_minutes": info["gpu_minutes"], "device": info["device"], | |
| "val": {k: vm[k]["value"] for k in keys if k in vm}, | |
| "test": {k: tm[k]["value"] for k in keys if k in tm}, | |
| "test_metrics_json": recs["test"]["metrics_path"], "val_metrics_json": recs["val"]["metrics_path"]}) | |
| fs = info.get("f1sel") | |
| if fs: | |
| fnotes = (f"{a.notes};DUAL_SELECTION:checkpoint chosen on val mean macro-F1 (epoch {fs['epoch']}, " | |
| f"val F1 {fs['val_f1_mean']:.2f}); geometry of this checkpoint reported alongside").strip(";") | |
| frecs = {} | |
| for split in ("val", "test"): | |
| frecs[split] = run_harness(fs["preds"][split], fs["name"], a.model, a.seed, split, fnotes, a.n_boot, refs=refs, | |
| write_csv=not a.no_csv) | |
| print(summary_line(frecs[split]), flush=True) | |
| ftm, fvm = frecs["test"]["metrics"], frecs["val"]["metrics"] | |
| write_run_json(Path(fs["out_dir"]), fs["name"], [__file__, P.__file__, M.__file__] + info["code_files"], [P.CACHE_NPZ], | |
| seed=a.seed, t_start=info["t_start"], input_listing=False, extra={ | |
| "config": vars(a), "selection": "val_mean_macro_f1 (dual selection)", "parent_run": name, | |
| "best_epoch": fs["epoch"], "best_val_f1_mean": fs["val_f1_mean"], | |
| "val": {k: fvm[k]["value"] for k in keys if k in fvm}, | |
| "test": {k: ftm[k]["value"] for k in keys if k in ftm}, | |
| "test_metrics_json": frecs["test"]["metrics_path"], "val_metrics_json": frecs["val"]["metrics_path"]}) | |
| pend = out_dir / "pending_eval.json" | |
| if pend.exists(): | |
| pend.rename(out_dir / "pending_eval.done.json") | |
| print(f"[{name}] done in {(time.time() - info['t_start']) / 60:.1f} min (GPU {info['gpu_minutes']:.1f} min); " | |
| f"{info['stop_reason']}", flush=True) | |
| return recs | |
| def finish_eval(out_dir): | |
| out_dir = Path(out_dir) | |
| pend = out_dir / "pending_eval.json" | |
| if not pend.exists(): | |
| print(f"[finish_eval] nothing pending in {out_dir}") | |
| return None | |
| return evaluate_and_record(json.loads(pend.read_text())) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 59.3 kB
- Xet hash:
- fb6b7778f9e317252b055db97794e74a6cf0f2f03daff87ac7e9c1da270d560a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.