Buckets:
| """cwm3/aide/train/pretrain_cwm3d_baton.py — self-supervised pretraining of CWM3-D-kp on the BATON in-cabin pose corpus (GPU-baselines A; C21). | |
| Corpus: aide/results/cache/baton_pose_corpus.npz (aide/features/build_baton_pose_corpus.py — TRAIN-split, train-only-device BATON routes; 10-frame | |
| windows at 0.3 s steps in the AIDE pixel convention; per-window isotropic shoulder-based similarity into the AIDE frame; per-part confidence quantile maps to the AIDE | |
| train score distribution). Frames are re-read from the route npz, transformed and confidence-mapped here, and kept on the GPU as one f16 | |
| tensor; every batch is a gather + the transform + the train_kp augmentation (scale ±5 %, translate ±20 px, keypoint noise 2(1-conf) px, joint | |
| dropout 0.1 on the observed frames — the same operations as aide/train/train_kp.TrainSampler). | |
| Task = the AIDE protocol task: 5 observed frames -> 5 future frames of the same 136 keypoints, loss = train_kp's (Euclidean + 0.1 Laplace NLL + | |
| 0.05 bone + 0.01 jerk). Model = cwm3d.build("cwm3d", d 256, 4 layers, 4 heads, ffn 1024, head 512) built right after torch.manual_seed(seed), i.e. | |
| the SAME initial weights as aide/results/runs/cwm3d_kp_s<seed>_v2 (train_kp seeds and builds in the same order). | |
| Selection / early stop: All-MPJPE (EMA weights) on the corpus' held-out routes (route_val; never fitted). The AIDE VAL split is evaluated | |
| zero-shot after every epoch as a DIAGNOSTIC curve only (not used for selection); AIDE test is never touched here except by the optional | |
| zero-shot run below. | |
| Outputs aide/results/runs/<name>/: ckpt_best.pt ({"model": EMA state dict, "epoch", "val": ..., "args"}), ckpt_last.pt, history.json, run json | |
| (ms/step, samples/s, GPU minutes, corpus digest). --zeroshot (default on): the best checkpoint predicts AIDE val + test ONCE as run | |
| `cwm3d_kp_batonpre_zeroshot_s<seed>_v2` (no AIDE selection at all) through the standard harness (paired Δ vs ZV / refs). | |
| Fine-tuning = aide/train/train_kp.py --init-ckpt <ckpt_best.pt> (opt-in flag, default off). | |
| --model motionbert_kp | stgcn_kp (arch attribution (1), decisions 2026-09-14): the SAME corpus / windows / augmentation / task / | |
| optimiser / schedule / EMA / selection with the keypoint baseline built right after torch.manual_seed(seed) (== train_kp's init for | |
| that seed); these models return a plain tensor, so LS.compute_losses reduces to the Euclidean term — exactly the loss of their AIDE | |
| recipe (train_kp --loss euc). Runs are named <model>_batonpre_pretrain_s<seed> / <model>_batonpre_zeroshot_s<seed>_v2. | |
| --out-root redirects the run directory (smoke tests). Default --model cwm3d: behaviour unchanged. | |
| --model driverwm_kp (P2, decisions 2026-09-15 "Driver-WM architecture on CockpitWM tokens + BATON pretraining"): aide/models/driverwm_kp.py built | |
| right after torch.manual_seed(seed) with train_kp's kwargs (dropout 0) == the init of driverwm_kp_s<seed>_v2. The BATON corpus has no exterior | |
| Qwen features: the model is called without `vis`, i.e. its own missing-exterior path (forward(obs, vis=None) -> zero exterior features of the | |
| right shape), so the exterior stream (view / step embeddings, pre_ext, f_ext, the context keys / values and the gate) runs on a constant zero | |
| input during pretraining — recorded as a deviation in the run json (`exterior_stream`). Tensor output -> Euclidean term only (its AIDE recipe | |
| is --loss euc). Fine-tuning: train_kp --init-parts trunk loads driverwm_kp.DriverWMKP.TRUNK_PREFIXES = everything except the zero-initialised | |
| residual read-out dec.out.* (the fine-tune starts at exact zero-velocity). Run name driverwm_kp_batonpre_pretrain_s<seed>. | |
| --motion-scale-max S (P1-ext variant `scale`, decisions 2026-09-15; default 1.0 = off, every other variant unaffected): per-window factor | |
| f ~ U[1, S] multiplies every frame's (x, y) displacement relative to the LAST OBSERVED pose (future and observed window consistently; the last | |
| observed pose and the confidences are untouched), applied to the training batches BEFORE the standard augmentation (rationale: BATON future | |
| displacement median 13.3 px vs AIDE 29.1 px). Function `motion_scale` (unit-tested: S = 1 is a no-op, factor 2 doubles the displacements). | |
| GPU etiquette (project rule): no lock; threads <= 4; launch gates = free VRAM >= need + 4 GB, MemAvailable >= 14 GB + own RSS, trainers alive | |
| (train_cwm3|train_kp|pretrain python processes) <= 2 before launch. No BATON label, CAN or DMS field is read (pose keypoints / boxes only). | |
| python aide/train/pretrain_cwm3d_baton.py --seed 42 [--epochs 6] [--bs 256] [--max-minutes 90] [--limit-windows N] [--cpu] | |
| """ | |
| from __future__ import annotations | |
| import os | |
| 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 json # noqa: E402 | |
| import math # noqa: E402 | |
| import re # 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.features.build_baton_pose_corpus import fill_missing_joints # noqa: E402 (HALPE 17 proxy, identical to the corpus builder) | |
| from aide.models import cwm3d, heads as HD, losses as LS # noqa: E402 | |
| from aide.train.train_kp import EMA, PROTO_IDX, seed_all, split_obs_fut # noqa: E402 | |
| from common.runmeta import sha256_file, write_run_json # noqa: E402 | |
| CORPUS = P.RESULTS / "cache" / "baton_pose_corpus.npz" | |
| RUNS = P.RESULTS / "runs" | |
| MIN_FREE_VRAM_GB, MIN_AVAIL_RAM_GB, MAX_OTHER_TRAINERS = 12.0, 20.0, 2 | |
| PART_SLICES = {k: (int(v[0]), int(v[-1]) + 1) for k, v in P.PARTS.items()} | |
| def log(msg, out): | |
| line = f"[{time.strftime('%H:%M:%S')}] {msg}" | |
| print(line, flush=True) | |
| with open(Path(out) / "train.log", "a") as f: | |
| f.write(line + "\n") | |
| def live_trainers(): | |
| me = os.getpid() | |
| n = [] | |
| for pid in os.listdir("/proc"): | |
| if not pid.isdigit() or int(pid) == me: | |
| continue | |
| try: | |
| cmd = open(f"/proc/{pid}/cmdline", "rb").read().replace(b"\0", b" ").decode(errors="ignore") | |
| except OSError: | |
| continue | |
| if re.search(r"train_cwm3|train_kp|pretrain", cmd) and "pgrep" not in cmd and re.match(r"\S*python[\d.]*\s", cmd): | |
| n.append((int(pid), cmd[:120])) | |
| return n | |
| def check_resources(out, cpu, min_free_vram_gb=MIN_FREE_VRAM_GB, min_avail_ram_gb=MIN_AVAIL_RAM_GB): | |
| avail_gb = float([l for l in open("/proc/meminfo") if l.startswith("MemAvailable")][0].split()[1]) / 2**20 | |
| alive = live_trainers() | |
| msg = f"resources: RAM available {avail_gb:.1f} GB, other trainers alive {len(alive)}" | |
| if not cpu: | |
| free_gb = torch.cuda.mem_get_info()[0] / 2**30 | |
| msg += f", VRAM free {free_gb:.1f} GB" | |
| assert free_gb >= min_free_vram_gb, f"free VRAM {free_gb:.1f} GB < {min_free_vram_gb}" | |
| log(msg, out) | |
| assert avail_gb >= min_avail_ram_gb, f"available RAM {avail_gb:.1f} GB < {min_avail_ram_gb}" | |
| assert len(alive) <= MAX_OTHER_TRAINERS, f"{len(alive)} trainers alive (limit {MAX_OTHER_TRAINERS} others): {[c for _, c in alive]}" | |
| def load_corpus(path, device, limit=0, seed=0, out=None, limit_routes=0): | |
| """corpus index -> GPU tensors: K [F,136,3] f16 (conf quantile-mapped), win_base [M] i64 (frame index of sample 0 in K), wparams [M,5] | |
| (cx, cy, scale, native_w, native_h), win_val [M] bool, meta dict.""" | |
| z = np.load(path, allow_pickle=True) | |
| stride = int(z["frame_stride"]) | |
| uids, srcs, wh, nfr, rval = z["route_uid"].astype(str), z["route_src"].astype(str), z["route_native_wh"], z["route_n_frames"].astype(np.int64), z["route_val"].astype(bool) | |
| q_b, q_a = z["conf_q_baton"].astype(np.float64), z["conf_q_aide"].astype(np.float64) | |
| params = json.loads(str(z["params"])) | |
| dirs = params["pose_dirs"] | |
| wr, ws = z["win_route"].astype(np.int64), z["win_start"].astype(np.int64) | |
| if limit_routes: # smoke / CPU tests: the first `limit_routes` routes only (host RAM stays small) | |
| keep_r = np.arange(min(limit_routes, len(uids))) | |
| sel_w = wr < len(keep_r) | |
| uids, srcs, wh, nfr, rval = uids[keep_r], srcs[keep_r], wh[keep_r], nfr[keep_r], rval[keep_r] | |
| wr, ws = wr[sel_w], ws[sel_w] | |
| z_cx, z_cy, z_sc = z["win_cx"][sel_w], z["win_cy"][sel_w], z["win_scale"][sel_w] | |
| else: | |
| z_cx, z_cy, z_sc = z["win_cx"], z["win_cy"], z["win_scale"] | |
| off = np.concatenate([[0], np.cumsum(nfr)]) | |
| F = int(off[-1]) | |
| K = np.empty((F, P.J, 3), np.float16) | |
| t0 = time.time() | |
| for i, (u, s) in enumerate(zip(uids, srcs)): | |
| with np.load(Path(dirs[s]) / f"{u}.npz") as r: | |
| k = fill_missing_joints(r["kpts"].astype(np.float32)) | |
| assert len(k) == nfr[i], (u, len(k), nfr[i]) | |
| c = np.nan_to_num(k[..., 2]) | |
| for j, (lo, hi) in enumerate(PART_SLICES.values()): | |
| xp = q_b[j] + np.arange(len(q_b[j])) * 1e-7 # strictly increasing for np.interp | |
| c[:, lo:hi] = np.interp(c[:, lo:hi], xp, q_a[j]) | |
| k[..., 2] = np.clip(c, 0.0, 1.0) | |
| K[off[i]:off[i + 1]] = np.nan_to_num(k).astype(np.float16) | |
| if out is not None and (i + 1) % 100 == 0: | |
| log(f" loaded {i + 1}/{len(uids)} routes ({time.time() - t0:.0f}s)", out) | |
| base = off[wr] + ws | |
| wp = np.stack([z_cx, z_cy, z_sc, wh[wr, 0].astype(np.float32), wh[wr, 1].astype(np.float32)], 1).astype(np.float32) | |
| wval = rval[wr] | |
| if limit and limit < len(base): | |
| sel = np.sort(np.random.default_rng(seed).choice(len(base), limit, replace=False)) | |
| base, wp, wval = base[sel], wp[sel], wval[sel] | |
| meta = {"corpus": str(path), "corpus_sha256": sha256_file(path), "n_routes": int(len(uids)), "n_routes_val": int(rval.sum()), "n_frames": F, "n_windows": int(len(base)), | |
| "n_windows_val": int(wval.sum()), "frame_stride": stride, "params": params, "stats": json.loads(str(z["stats"])), "aide_target": z["aide_target"].tolist()} | |
| T = {"K": torch.from_numpy(K).to(device), "base": torch.from_numpy(base).to(device), "wp": torch.from_numpy(wp).to(device), | |
| "offs": (stride * torch.arange(P.N_SAMPLE, device=device)).long(), "box_t": torch.as_tensor(z["aide_target"], device=device, dtype=torch.float32)} | |
| return T, torch.from_numpy(wval), meta | |
| def model_pred(out): | |
| return out["pred"] if isinstance(out, dict) else out | |
| def build_model(a): | |
| """-> (model, code files). cwm3d: unchanged; motionbert_kp / stgcn_kp: train_kp's build call (dropout 0) for the shared-recipe baselines.""" | |
| if a.model == "cwm3d": | |
| return cwm3d.build("cwm3d", d=a.d, layers=a.layers, heads=a.heads, ffn=a.ffn, head_hidden=a.head_hidden, dropout=0.0), [cwm3d.__file__, HD.__file__] | |
| if a.model == "cwm3d_v2": | |
| # K-v2 (plan A.1, 2026-09-14): the same build call as train_kp.build_model (kwargs_from_args; dropout 0) -> identical seed init | |
| from aide.models import cwm3d_spatial, cwm3d_v2, motionbert_kp | |
| kw = {**cwm3d_v2.kwargs_from_args(a), "dropout": 0.0} | |
| return cwm3d_v2.build("cwm3d_v2", **kw), [cwm3d_v2.__file__, HD.__file__, cwm3d_spatial.__file__, motionbert_kp.__file__, LS.__file__] | |
| if a.model == "motionbert_kp": | |
| from aide.models import motionbert_kp | |
| return motionbert_kp.build("motionbert_kp", dropout=0.0), [motionbert_kp.__file__] | |
| if a.model == "stgcn_kp": | |
| from aide.models import driverwm_repro, motionbert_kp, stgcn_kp | |
| return stgcn_kp.build("stgcn_kp", dropout=0.0), [stgcn_kp.__file__, driverwm_repro.__file__, motionbert_kp.__file__] | |
| if a.model == "driverwm_kp": | |
| # P2 (decisions 2026-09-15): train_kp.build_model's call for driverwm_kp (dropout 0 in the shared recipe) -> identical seed init | |
| from aide.models import driverwm_kp, driverwm_repro | |
| return driverwm_kp.build("driverwm_kp", dropout=0.0), [driverwm_kp.__file__, driverwm_repro.__file__, HD.__file__] | |
| raise ValueError(a.model) | |
| def gather(T, widx): | |
| """widx LongTensor [B] (device) -> frames [B,10,136,3] f32 in the AIDE pixel convention.""" | |
| fi = T["base"][widx][:, None] + T["offs"][None, :] | |
| x = T["K"][fi].float() | |
| p = T["wp"][widx] | |
| x[..., 0] = (x[..., 0] * p[:, 3, None, None] - p[:, 0, None, None]) * p[:, 2, None, None] + T["box_t"][0] | |
| x[..., 1] = (x[..., 1] * p[:, 4, None, None] - p[:, 1, None, None]) * p[:, 2, None, None] + T["box_t"][1] | |
| return x | |
| def augment(x, noise_px, joint_drop, translate, scale): | |
| """train_kp.TrainSampler.batch_extra's augmentation, verbatim in its operations (all 10 frames for the geometric part, observed frames for noise / dropout).""" | |
| B, device = x.shape[0], x.device | |
| if scale > 0: | |
| s = 1 + (torch.rand(B, 1, 1, 1, device=device) * 2 - 1) * scale | |
| c = x[:, P.T_OBS - 1, :, :2].mean(1)[:, None, None] | |
| x[..., :2] = c + (x[..., :2] - c) * s | |
| if translate > 0: | |
| x[..., :2] += (torch.rand(B, 1, 1, 2, device=device) * 2 - 1) * translate | |
| if noise_px > 0: | |
| n = noise_px * (1 - x[:, :P.T_OBS, :, 2:3]) * torch.randn_like(x[:, :P.T_OBS, :, :2]) | |
| x[:, :P.T_OBS, :, :2] += n | |
| if joint_drop > 0: | |
| drop = torch.rand(B, P.T_OBS, P.J, device=device) < 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]) | |
| return x | |
| def motion_scale(x, smax, factors=None): | |
| """P1-ext `scale` variant: x [B,10,J,3] (AIDE pixel convention) -> per-window factor f ~ U[1, smax] (or the given `factors` [B]) multiplies | |
| every frame's (x, y) displacement relative to the LAST OBSERVED pose (frame T_OBS - 1, the fine-tune anchor); that pose and the confidences | |
| are untouched. smax <= 1 with factors None -> x returned unchanged (exact no-op).""" | |
| if factors is None: | |
| if smax <= 1.0: | |
| return x | |
| factors = 1.0 + (smax - 1.0) * torch.rand(x.shape[0], device=x.device) | |
| f = torch.as_tensor(factors, dtype=x.dtype, device=x.device).reshape(-1, 1, 1, 1) | |
| last = x[:, P.T_OBS - 1:P.T_OBS, :, :2] | |
| x[..., :2] = last + (x[..., :2] - last) * f # frame T_OBS-1: last + 0 * f == last exactly | |
| return x | |
| def eval_corpus(model, T, widx_all, device, bs=1024): | |
| """held-out windows -> All-MPJPE, HM-MPJPE (top 10 % by GT motion), mean loss-free Euclidean per horizon.""" | |
| model.eval() | |
| errs, scores = [], [] | |
| for i in range(0, len(widx_all), bs): | |
| x = gather(T, widx_all[i:i + bs].to(device)) | |
| obs, fut = split_obs_fut(x) | |
| out = model(obs) | |
| e = torch.sqrt(((model_pred(out) - fut) ** 2).sum(-1)) # [B,5,J] | |
| errs.append(e.mean(2).cpu()); scores.append(torch.sqrt(((fut[:, 1:] - fut[:, :-1]) ** 2).sum(-1)).mean((1, 2)).cpu()) | |
| e = torch.cat(errs).numpy(); s = torch.cat(scores).numpy() | |
| top = np.argsort(-s, kind="stable")[:max(1, int(0.1 * len(s)))] | |
| return {"mpjpe_all": float(e.mean()), "mpjpe_hm": float(e[top].mean()), "per_horizon": e.mean(0).tolist(), "n": int(len(e)), "hm_n": int(len(top))} | |
| def predict_aide(model, S10, device, bs=256): | |
| model.eval() | |
| preds = [] | |
| for i in range(0, len(S10), bs): | |
| obs, _ = split_obs_fut(torch.as_tensor(S10[i:i + bs]).to(device).float()) | |
| preds.append(model_pred(model(obs)).float().cpu().numpy()) | |
| return np.concatenate(preds) | |
| def main(argv=None): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--seed", type=int, default=42) | |
| ap.add_argument("--model", default="cwm3d", choices=("cwm3d", "cwm3d_v2", "motionbert_kp", "stgcn_kp", "driverwm_kp")) | |
| ap.add_argument("--out-root", default=None, help="run dir root (default aide/results/runs)") | |
| ap.add_argument("--name", default=None) | |
| ap.add_argument("--corpus", default=str(CORPUS)) | |
| ap.add_argument("--epochs", type=int, default=6) | |
| ap.add_argument("--bs", type=int, default=256) | |
| ap.add_argument("--lr", type=float, default=3e-4) | |
| ap.add_argument("--wd", type=float, default=0.05) | |
| ap.add_argument("--warmup-epochs", type=float, default=1.0) | |
| ap.add_argument("--ema", type=float, default=0.999) | |
| ap.add_argument("--clip", type=float, default=1.0) | |
| ap.add_argument("--patience", type=int, default=3) | |
| 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("--noise-px", type=float, default=2.0) | |
| ap.add_argument("--joint-drop", type=float, default=0.1) | |
| ap.add_argument("--translate", type=float, default=20.0) | |
| ap.add_argument("--scale", type=float, default=0.05) | |
| ap.add_argument("--motion-scale-max", type=float, default=1.0, help="P1-ext `scale` variant: per-window motion factor ~ U[1, S] about the last observed pose (1.0 = off)") | |
| 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("--max-minutes", type=float, default=90.0) | |
| ap.add_argument("--limit-windows", type=int, default=0, help="seeded subsample of the corpus windows (smoke / budget)") | |
| ap.add_argument("--limit-routes", type=int, default=0, help="smoke: load only the first N routes of the corpus (0 = all)") | |
| ap.add_argument("--val-windows", type=int, default=20000, help="max held-out windows scored per epoch (seeded subsample)") | |
| ap.add_argument("--zeroshot", type=int, default=1, help="1: predict AIDE val + test once with the best checkpoint as run cwm3d_kp_batonpre_zeroshot_s<seed>_v2") | |
| ap.add_argument("--refs", default="zvshrink=zv_shrink_val,simlpe=simlpe_kp_s42_v2,kp=cwm3d_kp_s42_v2") | |
| ap.add_argument("--n-boot", type=int, default=10000) | |
| ap.add_argument("--cpu", action="store_true") | |
| ap.add_argument("--no-check", action="store_true", help="skip the resource gates (smoke on CPU)") | |
| ap.add_argument("--min-free-vram-gb", type=float, default=MIN_FREE_VRAM_GB, | |
| help="free-VRAM floor of the resource gate (default 12 = unchanged); the lane runner passes its need-based value (measured peak footprint + 4 GB)") | |
| ap.add_argument("--min-avail-ram-gb", type=float, default=MIN_AVAIL_RAM_GB, | |
| help="MemAvailable floor of the resource gate (default 20 = unchanged); the lane runner passes its need-based value (measured peak RSS + 10 GB)") | |
| from aide.models import cwm3d_v2 | |
| cwm3d_v2.add_args(ap) # K-v2 flags (--v2-*, --haz-*); inert for the other models | |
| a = ap.parse_args(argv) | |
| tag = "cwm3d_kp" if a.model == "cwm3d" else a.model | |
| name = a.name or f"{tag}_batonpre_pretrain_s{a.seed}" | |
| runs_root = Path(a.out_root) if a.out_root else RUNS | |
| out = runs_root / name | |
| out.mkdir(parents=True, exist_ok=True) | |
| t_start = time.time() | |
| device = torch.device("cpu" if a.cpu else "cuda") | |
| if not a.no_check: | |
| check_resources(out, a.cpu, a.min_free_vram_gb, a.min_avail_ram_gb) | |
| log(f"=== pretrain_cwm3d_baton {name}: {vars(a)}", out) | |
| seed_all(a.seed) | |
| model, model_files = build_model(a) # built right after seed_all == train_kp's init for this seed | |
| model = model.to(device) | |
| n_params = sum(p.numel() for p in model.parameters()) | |
| T, wval, meta = load_corpus(a.corpus, device, a.limit_windows, a.seed, out, a.limit_routes) | |
| tr_idx = torch.nonzero(~wval).flatten() | |
| va_idx = torch.nonzero(wval).flatten() | |
| if len(va_idx) == 0: # smoke subsets without a held-out route: last 10 % of the windows | |
| n_va = max(1, len(tr_idx) // 10) | |
| va_idx, tr_idx = tr_idx[-n_va:], tr_idx[:-n_va] | |
| if len(va_idx) > a.val_windows: | |
| va_idx = va_idx[torch.as_tensor(np.sort(np.random.default_rng(a.seed).choice(len(va_idx), a.val_windows, replace=False)))] | |
| log(f"corpus: {meta['n_routes']} routes ({meta['n_routes_val']} held out), {meta['n_frames']:,} frames, windows fit {len(tr_idx):,} / val {len(va_idx):,} " | |
| f"(GPU tensor {T['K'].numel() * 2 / 2**30:.2f} GiB); model params {n_params:,}", out) | |
| c = np.load(P.CACHE_NPZ, allow_pickle=False) | |
| va = c["split"] == "val" | |
| S10_va, id_va = c["S10"][va], c["clip_id"][va] | |
| fut_va = S10_va[:, P.T_OBS:] | |
| hm_va, _, _ = M.hm_mask(fut_va) | |
| haz_tau, hazard_info = None, None | |
| if getattr(model, "haz", None) is not None: | |
| # K-v2: displacement thresholds / cumulative onset prior from the AIDE TRAIN split (the BATON corpus carries no labels; poses only), | |
| # written into the model BEFORE the EMA copy; the hazard BCE (weight --haz-weight) is added to the pretraining loss below | |
| tau, prior = LS.hazard_stats(c["S"][c["split"] == "train"][:, PROTO_IDX], a.haz_quantile) | |
| model.set_hazard_prior(tau, prior) | |
| haz_tau = torch.as_tensor(tau, dtype=torch.float32, device=device) | |
| hazard_info = {"tau": tau.tolist(), "prior_P": prior.tolist(), "quantile": float(a.haz_quantile), "weight": float(a.haz_weight), | |
| "source": "AIDE train split, S[:, PROTO_IDX]"} | |
| log(f"hazard prior (AIDE train, q={a.haz_quantile}): tau {np.round(tau, 1).tolist()} px, P_h5 {np.round(prior[:, -1], 3).tolist()}", out) | |
| ema = EMA(model, a.ema) | |
| opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.wd) | |
| steps_per_epoch = math.ceil(len(tr_idx) / a.bs) | |
| total = a.epochs * steps_per_epoch | |
| warm = max(1, int(a.warmup_epochs * 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)))) | |
| w = {"euc": a.w_euc, "nll": a.w_nll, "bone": a.w_bone, "jerk": a.w_jerk} | |
| g = torch.Generator(device="cpu").manual_seed(a.seed) | |
| # epoch-0 diagnostics (untrained = zero-velocity by construction) | |
| v0 = eval_corpus(ema.shadow, T, va_idx, device) | |
| pv0 = predict_aide(ema.shadow, S10_va, device) | |
| e0 = M.joint_errors(pv0, fut_va) | |
| log(f"ep -1 (init = zero-velocity): corpus val All {v0['mpjpe_all']:.2f} HM {v0['mpjpe_hm']:.2f} | AIDE val zero-shot All {M.mpjpe(e0):.3f} HM {M.mpjpe(e0[hm_va]):.2f}", out) | |
| hist, best, best_ep, bad, stop = [], float("inf"), -1, 0, "epochs" | |
| t_gpu = time.time() | |
| for ep in range(a.epochs): | |
| model.train() | |
| perm = tr_idx[torch.randperm(len(tr_idx), generator=g)] | |
| acc, nb = {}, 0 | |
| if device.type == "cuda": | |
| torch.cuda.synchronize() | |
| t_ep = time.time() | |
| for i in range(0, len(perm), a.bs): | |
| x = gather(T, perm[i:i + a.bs].to(device)) | |
| x = motion_scale(x, a.motion_scale_max) # no-op unless --motion-scale-max > 1 | |
| x = augment(x, a.noise_px, a.joint_drop, a.translate, a.scale) | |
| obs, fut = split_obs_fut(x) | |
| out_ = model(obs) | |
| loss, terms = LS.compute_losses(out_, fut, w, "euc") | |
| if haz_tau is not None and a.haz_weight > 0 and "haz_logits" in out_: # K-v2 onset-hazard BCE on the corpus future block | |
| on, rk = LS.hazard_targets(obs[:, -1, :, :2], fut, haz_tau) | |
| l_haz = LS.hazard_bce(out_["haz_logits"], on, rk) | |
| loss = loss + a.haz_weight * l_haz | |
| terms = {**terms, "haz": l_haz.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) + float(loss.detach()) | |
| for k, v in terms.items(): | |
| acc[k] = acc.get(k, 0.0) + float(v) | |
| nb += 1 | |
| if nb % 500 == 0: | |
| log(f" ep {ep} step {nb}/{steps_per_epoch} loss {acc['total'] / nb:.3f} ({(time.time() - t_ep) / nb * 1000:.1f} ms/step)", out) | |
| if (time.time() - t_gpu) / 60 > a.max_minutes: | |
| break | |
| if device.type == "cuda": | |
| torch.cuda.synchronize() | |
| ms_step = (time.time() - t_ep) / max(nb, 1) * 1000 | |
| vc = eval_corpus(ema.shadow, T, va_idx, device) | |
| pv = predict_aide(ema.shadow, S10_va, device) | |
| e = M.joint_errors(pv, fut_va) | |
| rec = {"epoch": ep, "steps": nb, "ms_per_step": round(ms_step, 2), "samples_per_s": round(a.bs * 1000 / max(ms_step, 1e-6), 1), "lr": sched.get_last_lr()[0], | |
| "corpus_val_all": vc["mpjpe_all"], "corpus_val_hm": vc["mpjpe_hm"], "corpus_val_per_horizon": vc["per_horizon"], | |
| "aide_val_zeroshot_all": M.mpjpe(e), "aide_val_zeroshot_hm": M.mpjpe(e[hm_va]), "t_min": round((time.time() - t_start) / 60, 2), | |
| **{f"train_{k}": v / max(nb, 1) for k, v in acc.items()}} | |
| hist.append(rec) | |
| improved = vc["mpjpe_all"] < best - 1e-3 | |
| if improved: | |
| best, best_ep, bad = vc["mpjpe_all"], ep, 0 | |
| torch.save({"model": ema.shadow.state_dict(), "epoch": ep, "val": vc, "aide_val_zeroshot": {"all": rec["aide_val_zeroshot_all"], "hm": rec["aide_val_zeroshot_hm"]}, | |
| "args": vars(a), "seed": a.seed, "corpus_sha256": meta["corpus_sha256"], "n_params": n_params, "source": "aide/train/pretrain_cwm3d_baton.py"}, out / "ckpt_best.pt") | |
| else: | |
| bad += 1 | |
| torch.save({"model": ema.shadow.state_dict(), "epoch": ep, "args": vars(a)}, out / "ckpt_last.pt") | |
| log(f"ep {ep}: loss {acc['total'] / max(nb, 1):.3f} [" + " ".join(f"{k} {v / max(nb, 1):.3f}" for k, v in acc.items() if k != "total") + f"] {ms_step:.1f} ms/step " | |
| f"| corpus val All {vc['mpjpe_all']:.2f} HM {vc['mpjpe_hm']:.2f} best {best:.2f}@{best_ep} | AIDE val zero-shot All {rec['aide_val_zeroshot_all']:.3f} HM {rec['aide_val_zeroshot_hm']:.2f} " | |
| f"| {rec['t_min']:.1f} min", out) | |
| (out / "history.json").write_text(json.dumps(hist, indent=1)) | |
| if (time.time() - t_gpu) / 60 > a.max_minutes: | |
| stop = f"time_budget({a.max_minutes}min)" | |
| break | |
| if bad >= a.patience: | |
| stop = f"early_stop(patience={a.patience})" | |
| break | |
| gpu_min = (time.time() - t_gpu) / 60 | |
| ck = torch.load(out / "ckpt_best.pt", map_location=device) | |
| ema.shadow.load_state_dict(ck["model"]) | |
| (out / "TRAIN_DONE").write_text(json.dumps({"best_epoch": best_ep, "best_corpus_val_all": best, "epochs_run": len(hist), "stop": stop})) | |
| zs = None | |
| if a.zeroshot: | |
| zname = f"{tag}_batonpre_zeroshot_s{a.seed}_v2" | |
| zdir = runs_root / zname | |
| zdir.mkdir(parents=True, exist_ok=True) | |
| te = c["split"] == "test" | |
| zs = {"name": zname, "preds": {}} | |
| for split, S10, ids in (("val", S10_va, id_va), ("test", c["S10"][te], c["clip_id"][te])): | |
| pr = predict_aide(ema.shadow, S10, device) | |
| assert np.isfinite(pr).all() | |
| np.savez(zdir / f"pred_{split}.npz", pred=pr.astype(np.float32), clip_id=ids.astype(str), split=np.array(split)) | |
| zs["preds"][split] = str(zdir / f"pred_{split}.npz") | |
| refs = parse_refs([r for r in a.refs.split(",") if r]) | |
| recs = {} | |
| for split in ("val", "test"): | |
| recs[split] = run_harness(zs["preds"][split], zname, a.model, a.seed, split, f"baton_pretrain_zeroshot;pretrain_run={name};best_ep={best_ep}", a.n_boot, refs=refs) | |
| log(summary_line(recs[split]), out) | |
| 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_kp_all", "diff_vs_kp_hm") | |
| zs.update(val={k: recs["val"]["metrics"][k]["value"] for k in keys if k in recs["val"]["metrics"]}, test={k: recs["test"]["metrics"][k]["value"] for k in keys if k in recs["test"]["metrics"]}) | |
| write_run_json(zdir, zname, [Path(__file__), *model_files, LS.__file__, P.__file__, M.__file__], [P.CACHE_NPZ, Path(a.corpus)], seed=a.seed, t_start=t_start, input_listing=False, | |
| extra={"config": vars(a), "selection": "NONE on AIDE (BATON held-out routes only; zero-shot)", "pretrain_run": name, "best_epoch": best_ep, "n_params": n_params, | |
| "val": zs["val"], "test": zs["test"], "test_metrics_json": recs["test"]["metrics_path"], "val_metrics_json": recs["val"]["metrics_path"], "dms_fields_used": []}) | |
| write_run_json(out, name, [Path(__file__), *model_files, LS.__file__, P.__file__, M.__file__], [P.CACHE_NPZ, Path(a.corpus)], seed=a.seed, t_start=t_start, input_listing=False, | |
| extra={"config": vars(a), "n_params": n_params, "corpus": meta, "best_epoch": best_ep, "best_corpus_val_all": best, "epochs_run": len(hist), "stop_reason": stop, | |
| "gpu_minutes": round(gpu_min, 2), "ms_per_step_mean": float(np.mean([h["ms_per_step"] for h in hist])) if hist else None, "history": hist, "zeroshot": zs, | |
| "device": str(device), "dms_fields_used": [], "model": a.model, "loss_terms_active": sorted(k for k in acc if k != "total"), | |
| "hazard": hazard_info, "motion_scale_max": float(a.motion_scale_max), | |
| "exterior_stream": ("zero input: forward(obs, vis=None) -> the model's missing-exterior path (BATON corpus has no exterior Qwen features); " | |
| "deviation from the AIDE runs, which pass the exterior features") if a.model == "driverwm_kp" else None, | |
| "trunk_prefixes": list(getattr(model, "TRUNK_PREFIXES", ("tok.", "blocks.", "ln_f."))), | |
| "init_note": f"torch.manual_seed(seed) -> {a.model} build: same initial weights as {tag}_s<seed>_v2"}) | |
| log(f"done: best epoch {best_ep} corpus val All {best:.2f}; GPU {gpu_min:.1f} min; {stop}; ckpt {out / 'ckpt_best.pt'}", out) | |
| print("PRETRAIN_BATON_DONE", flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 30.8 kB
- Xet hash:
- 16344dc9d0b4ede8ffa40814f632f4ef6bd4c75667567e6967bdabb6ca641199
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.