Buckets:
| """cwm3/aide/train/pretrain_aide_multistride.py — F7 (a): AIDE-ONLY multi-stride self-supervised pretraining of the CockpitWM keypoint | |
| forecaster (docs/decisions.md F7 pre-registration row, 2026-09-17). ISOLATED from BATON: the only data read is the AIDE keypoint cache | |
| (aide/results/cache/aide_kp.npz, TRAIN split clips; val clips for selection only); no BATON route, corpus, target or hold-out path is opened. | |
| Task = the AIDE protocol task (5 observed -> 5 future frames of the 136 HALPE keypoints, self-supervised: no label is read), but the 10-frame | |
| windows are drawn from the 45-frame TRAIN clips at MULTIPLE strides: uniform stride g in --strides (1..4 frames = 15/7.5/5/3.75 Hz) plus the | |
| train_kp `mixed` family (gaps in {4, 5}, the protocol grid's family); per epoch every train clip contributes --windows-per-clip windows | |
| (stride and start uniform). Augmentation = train_kp.TrainSampler verbatim (scale +-5 %, translate +-20 px, keypoint noise 2(1-conf) px, joint | |
| dropout 0.1 on the observed frames). Loss = train_kp's (Euclidean + 0.1 Laplace NLL + 0.05 bone + 0.01 jerk; tensor-output models -> Euclidean). | |
| Model built right after seed_all(seed) with train_kp's kwargs (== the init of <model>_s<seed>_v2); --init-ckpt (optional, all matching tensors) | |
| starts the pretraining from an existing checkpoint (arm a2: the adopted BATON pretraining, whose corpus touches none of the 60 hold-out devices). | |
| Selection / early stop: AIDE VAL All-MPJPE of the EMA weights on the protocol frames (the validation split; never test — AIDE test is not read | |
| here at all). Outputs aide/results/runs/<name>/: ckpt_best.pt ({"model": EMA state dict, "epoch", "val", "args"}), ckpt_last.pt, history.json, | |
| train.log, TRAIN_DONE, run json. Fine-tuning = train_kp --init-ckpt <ckpt_best.pt> --init-parts trunk (adopted recipe verbatim). | |
| GPU etiquette: threads <= 4, no lock; the launch gate is applied by the lane runner (aide/train/run_f7_pretrain_distil.sh). | |
| python aide/train/pretrain_aide_multistride.py --seed 123 --name cwm3d_kp_aidepre_pretrain_s123 | |
| """ | |
| 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 json # noqa: E402 | |
| import math # 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.models import losses as LS # noqa: E402 | |
| from aide.train.train_kp import EMA, TrainSampler, build_model, load_init_ckpt, model_pred, predict, random_grids, seed_all, split_obs_fut # noqa: E402 | |
| from common.runmeta import sha256_file, write_run_json # noqa: E402 | |
| RUNS = P.RESULTS / "runs" | |
| FORBIDDEN_PATH_TOKENS = ("baton_pose_corpus", "pose_causal_routes", "results/holdout", "/data/backup/BATON", "canonical", "route_") | |
| def assert_aide_only(argv, holdout_file=ROOT / "results" / "holdout" / "holdout_leaves.txt"): | |
| """F7 isolation gate: no argument may name a BATON corpus / route / hold-out path or any hold-out device id.""" | |
| joined = " ".join(str(x) for x in argv) | |
| for tok in FORBIDDEN_PATH_TOKENS: | |
| assert tok not in joined, f"F7 isolation: forbidden path token {tok!r} in argv" | |
| if Path(holdout_file).exists(): | |
| for line in Path(holdout_file).read_text().splitlines(): | |
| parts = line.strip().split("/") | |
| if len(parts) >= 2: | |
| assert parts[1] not in joined, f"F7 isolation: hold-out device {parts[1]} referenced in argv" | |
| return True | |
| def parse_strides(s: str): | |
| out = [] | |
| for x in [t for t in s.split(",") if t]: | |
| out.append("mixed" if x == "mixed" else int(x)) | |
| if out[-1] != "mixed": | |
| assert 1 <= out[-1] <= (P.N_FRAMES_STD - 1) // (P.N_SAMPLE - 1), f"stride {out[-1]} does not fit a {P.N_FRAMES_STD}-frame clip" | |
| assert out, "empty --strides" | |
| return out | |
| def multistride_grids(rng: np.random.Generator, B: int, strides, n=P.N_FRAMES_STD): | |
| """[B,10] frame indices: per row a stride drawn uniformly from `strides` (ints = uniform gap, 'mixed' = train_kp.random_grids family | |
| with gaps in {4,5}); start uniform over the admissible range. Every index lies in [0, n-1].""" | |
| choice = rng.integers(0, len(strides), size=B) | |
| idx = np.zeros((B, P.N_SAMPLE), dtype=np.int64) | |
| for k, st in enumerate(strides): | |
| rows = np.nonzero(choice == k)[0] | |
| if not len(rows): | |
| continue | |
| if st == "mixed": | |
| idx[rows] = random_grids(rng, len(rows), p_proto=0.0, n=n) | |
| else: | |
| span = st * (P.N_SAMPLE - 1) | |
| start = rng.integers(0, n - span, size=len(rows)) | |
| idx[rows] = start[:, None] + st * np.arange(P.N_SAMPLE)[None, :] | |
| assert idx.min() >= 0 and idx.max() <= n - 1 | |
| return idx | |
| class MultiStrideSampler(TrainSampler): | |
| """TrainSampler whose grids come from multistride_grids (augmentation code inherited verbatim).""" | |
| def __init__(self, S45, seed, strides, **kw): | |
| super().__init__(S45, aug_window=True, seed=seed, **kw) | |
| self.strides = strides | |
| def batch_extra(self, idx, device): | |
| B = len(idx) | |
| grids = multistride_grids(self.rng, B, self.strides) | |
| ti = torch.as_tensor(idx, device=self.S.device) | |
| x = self.S[ti[:, None], torch.as_tensor(grids, device=self.S.device)].to(device, non_blocking=True).float().clone() | |
| # the geometric / noise / dropout augmentation below is TrainSampler.batch_extra's, applied to the multi-stride window | |
| 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] | |
| 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.noise_px > 0: | |
| x[:, :P.T_OBS, :, :2] += self.noise_px * (1 - x[:, :P.T_OBS, :, 2:3]) * torch.randn_like(x[:, :P.T_OBS, :, :2]) | |
| 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]) | |
| obs, fut = split_obs_fut(x) | |
| return obs, fut, {"grids": grids} | |
| def get_args(argv=None): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--seed", type=int, default=123) | |
| ap.add_argument("--model", default="cwm3d", choices=("cwm3d", "motionbert_kp")) | |
| ap.add_argument("--name", default=None) | |
| ap.add_argument("--out-root", default=None, help="run dir root (default aide/results/runs; smoke tests redirect)") | |
| ap.add_argument("--strides", default="1,2,3,4,mixed") | |
| ap.add_argument("--windows-per-clip", type=int, default=8, help="windows drawn per TRAIN clip per epoch") | |
| ap.add_argument("--epochs", type=int, default=60) | |
| ap.add_argument("--patience", type=int, default=10) | |
| ap.add_argument("--bs", type=int, default=128) | |
| 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("--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) | |
| # cwm3d architecture flags (train_kp defaults) | |
| 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("--dropout", type=float, default=0.0) | |
| ap.add_argument("--init-ckpt", default=None, help="optional: start from this checkpoint (all matching tensors; arm a2)") | |
| ap.add_argument("--init-parts", default="all", choices=("all", "trunk")) | |
| ap.add_argument("--max-minutes", type=float, default=30.0) | |
| ap.add_argument("--limit-train", type=int, default=0, help="smoke: first N train clips only") | |
| ap.add_argument("--cpu", action="store_true") | |
| ap.add_argument("--notes", default="") | |
| a = ap.parse_args(argv) | |
| # train_kp.build_model reads these attributes for cwm3d; keep them at the recipe defaults | |
| for k, v in dict(no_anchor=False, no_uncertainty=False, dct=False, rollout=None, visual=False, semantic=False, egomotion=False, | |
| refine_late=False, horizon_gain="none", formulation="residual").items(): | |
| setattr(a, k, v) | |
| return a | |
| def main(argv=None): | |
| argv = list(sys.argv[1:] if argv is None else argv) | |
| assert_aide_only(argv) | |
| a = get_args(argv) | |
| tag = "cwm3d_kp" if a.model == "cwm3d" else a.model | |
| name = a.name or f"{tag}_aidepre_pretrain_s{a.seed}" | |
| out = (Path(a.out_root) if a.out_root else RUNS) / name | |
| out.mkdir(parents=True, exist_ok=True) | |
| logf = open(out / "train.log", "a") | |
| def log(msg): | |
| line = f"[{time.strftime('%H:%M:%S')}] {msg}" | |
| print(line, flush=True); logf.write(line + "\n"); logf.flush() | |
| t_start = time.time() | |
| device = torch.device("cpu" if a.cpu else "cuda") | |
| log(f"=== pretrain_aide_multistride {name}: {vars(a)}") | |
| seed_all(a.seed) | |
| model, code_files = build_model(a.model, a) # built right after seed_all == train_kp's init for this seed | |
| model = model.to(device) | |
| init_report = None | |
| if a.init_ckpt: | |
| init_report = load_init_ckpt(model, a.init_ckpt, a.init_parts, device) | |
| log(f"init from {a.init_ckpt} ({a.init_parts}): {init_report['n_loaded']} tensors loaded, {init_report['n_missing_or_kept']} kept") | |
| n_params = sum(p.numel() for p in model.parameters()) | |
| c = np.load(P.CACHE_NPZ, allow_pickle=False) | |
| tr, va = (c["split"] == s for s in ("train", "val")) | |
| S45_tr = c["S"][tr] | |
| if a.limit_train: | |
| S45_tr = S45_tr[:a.limit_train] | |
| S10_va = c["S10"][va] | |
| fut_va = S10_va[:, P.T_OBS:] | |
| hm_va, _, _ = M.hm_mask(fut_va) | |
| strides = parse_strides(a.strides) | |
| sampler = MultiStrideSampler(S45_tr, a.seed, strides, noise_px=a.noise_px, joint_drop=a.joint_drop, translate=a.translate, scale=a.scale, device=device) | |
| ema = EMA(model, a.ema) | |
| opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.wd) | |
| n_win = len(S45_tr) * a.windows_per_clip | |
| steps_per_epoch = math.ceil(n_win / 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} | |
| rng = np.random.default_rng(a.seed) | |
| pv0 = predict(ema.shadow, S10_va, device) | |
| e0 = M.joint_errors(pv0, fut_va) | |
| log(f"train clips {len(S45_tr)} x {a.windows_per_clip} windows/epoch = {n_win} ({steps_per_epoch} steps); strides {strides}; params {n_params:,}; " | |
| f"AIDE val at init All {M.mpjpe(e0):.3f} HM {M.mpjpe(e0[hm_va]):.2f}") | |
| hist, best, best_ep, bad, stop = [], float("inf"), -1, 0, "epochs" | |
| t_gpu = time.time() | |
| for ep in range(a.epochs): | |
| model.train() | |
| perm = np.concatenate([rng.permutation(len(S45_tr)) for _ in range(a.windows_per_clip)]) | |
| acc, nb = {}, 0 | |
| t_ep = time.time() | |
| for i in range(0, len(perm), a.bs): | |
| obs, fut, _ = sampler.batch_extra(perm[i:i + a.bs], device) | |
| out_ = model(obs) | |
| loss, terms = LS.compute_losses(out_, fut, w, "euc") | |
| 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 (time.time() - t_gpu) / 60 > a.max_minutes: | |
| break | |
| ms_step = (time.time() - t_ep) / max(nb, 1) * 1000 | |
| pv = predict(ema.shadow, S10_va, device) | |
| e = M.joint_errors(pv, fut_va) | |
| v_all, v_hm = M.mpjpe(e), M.mpjpe(e[hm_va]) | |
| rec = {"epoch": ep, "steps": nb, "ms_per_step": round(ms_step, 2), "lr": sched.get_last_lr()[0], "val_mpjpe_all": v_all, "val_mpjpe_hm": v_hm, | |
| "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 = v_all < best - 1e-4 | |
| if improved: | |
| best, best_ep, bad = v_all, ep, 0 | |
| torch.save({"model": ema.shadow.state_dict(), "epoch": ep, "val": {"mpjpe_all": v_all, "mpjpe_hm": v_hm}, "args": vars(a), "seed": a.seed, | |
| "n_params": n_params, "source": "aide/train/pretrain_aide_multistride.py", "data": "AIDE train split only (aide_kp.npz)"}, 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.get('total', 0) / 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 | AIDE val All {v_all:.3f} HM {v_hm:.2f} best {best:.3f}@{best_ep} | {rec['t_min']:.1f} min") | |
| (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 | |
| (out / "TRAIN_DONE").write_text(json.dumps({"best_epoch": best_ep, "best_val_all": best, "epochs_run": len(hist), "stop": stop})) | |
| write_run_json(out, name, [Path(__file__), *code_files, LS.__file__, P.__file__, M.__file__], [P.CACHE_NPZ], seed=a.seed, t_start=t_start, input_listing=False, | |
| extra={"config": vars(a), "n_params": n_params, "best_epoch": best_ep, "best_val_mpjpe_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, | |
| "device": str(device), "dms_fields_used": [], "data": {"cache": str(P.CACHE_NPZ), "cache_sha256": sha256_file(P.CACHE_NPZ)[:16], "split": "train", | |
| "n_clips": int(len(S45_tr)), "strides": [str(s) for s in strides], "windows_per_clip": a.windows_per_clip}, | |
| "selection": "AIDE VAL All-MPJPE (EMA weights, protocol frames); AIDE test never read", "init": init_report, | |
| "isolation": "AIDE-only: no BATON route / corpus / target / hold-out path opened (assert_aide_only on argv)", | |
| "trunk_prefixes": list(getattr(model, "TRUNK_PREFIXES", ("tok.", "blocks.", "ln_f.")))}) | |
| log(f"done: best epoch {best_ep} AIDE val All {best:.3f}; GPU {gpu_min:.1f} min; {stop}; ckpt {out / 'ckpt_best.pt'}") | |
| print("PRETRAIN_AIDE_DONE", flush=True) | |
| logf.close() | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 16.2 kB
- Xet hash:
- 25c2cc2b0d38af4cd076242c9e844bc42458c40843f8593db7fd2f751c7e38d9
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.