Buckets:
| """Trainer for the faithful MotionBERT baseline (aide/models/motionbert_official) on the LOCKED AIDE keypoint harness. | |
| python aide/train/train_motionbert_official.py --formulation residual --init released --seed 42 --name mb_official_res_s42_v2 --no-lock | |
| python aide/train/train_motionbert_official.py --formulation abs --init released --seed 42 --name mb_official_abs_s42_v2 --no-lock | |
| python aide/train/train_motionbert_official.py --formulation residual --init random --seed 42 --name mb_official_rand_s42_v2 --no-lock | |
| Pre-registration docs/decisions.md 2026-09-16T21:28:48Z ('faithful MotionBERT baseline'). Everything that defines the protocol is imported from | |
| aide/train/train_kp.py and used unchanged: the cache (S [N,45,136,3] train windows / S10 val + test protocol frames), the split masks, TrainSampler | |
| (sub-window grids with p = 0.5 protocol grid, keypoint noise 2(1 - conf) px, joint dropout 0.1, translate +-20 px, scale +-5 %, no flip — the | |
| hardening-D recipe of motionbert_kp_s*_v2), split_obs_fut (the ONLY cut point: the model receives frames[:T_OBS]), predict (T_OBS frames only), | |
| EMA 0.999 weights for selection and prediction, early stopping on VAL All-MPJPE, plain Euclidean loss (aide/models/losses.compute_losses, | |
| --loss euc), ckpt_best / history.json / pred_{val,test}.npz outputs and evaluate_and_record (harness rows val + test with paired Δ vs | |
| zero-velocity / --refs, run_<name>_<utc>.json). The pre-registered differences: model = MotionBERTOfficial (official DSTformer, released | |
| weights or random init, residual or absolute read-out), AdamW with TWO parameter groups (backbone --lr-backbone 1e-4, read-out --lr-head 1e-3), | |
| weight decay 0.01, <= 60 epochs, patience 10 (10 warm-up epochs + cosine over the 60-epoch cap as in the harness), batch 32 (pushed through the | |
| 42.5 M-parameter backbone in --micro-bs slices with exact gradient accumulation: one augmentation draw, one clipped optimiser step, one EMA / scheduler | |
| step per 32-clip batch, identical to a full-batch step up to floating-point summation order — a GPU-memory device that keeps ~ 25 GB of VRAM free for | |
| the concurrent holdout extraction). Test is evaluated | |
| ONCE per run by the harness; nothing is selected on test. Peak GPU memory is recorded (`peak_gpu_gib`) for the queue's launch gate. | |
| """ | |
| 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 fcntl # 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.models.motionbert_official import adapter as A # noqa: E402 | |
| from aide.train.train_kp import EMA, GPU_LOCK, TrainSampler, evaluate_and_record, finish_eval, predict, seed_all # noqa: E402 | |
| MODEL_NAME = {("residual", "released"): "mb_official_res", ("abs", "released"): "mb_official_abs", ("residual", "random"): "mb_official_rand", | |
| ("abs", "random"): "mb_official_abs_rand"} | |
| DEFAULT_REFS = ["zvshrink=zv_shrink_val", "simlpe=simlpe_kp_s{seed}_v2", "kp=cwm3d_kp_s{seed}_v2", | |
| "kppre=cwm3d_kp_batonpre_trunkp123_s{seed}_v2", "mb=motionbert_kp_s{seed}_v2"] | |
| def get_args(argv=None): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--formulation", default="residual", choices=("residual", "abs")) | |
| ap.add_argument("--init", default="released", choices=("released", "random")) | |
| ap.add_argument("--ckpt", default="MB_release", choices=tuple(A.CKPTS)) | |
| ap.add_argument("--seed", type=int, default=42) | |
| ap.add_argument("--name", default=None) | |
| # pre-registered optimisation (the rest = harness recipe of motionbert_kp_s*_v2) | |
| ap.add_argument("--epochs", type=int, default=60) | |
| ap.add_argument("--patience", type=int, default=10) | |
| ap.add_argument("--bs", type=int, default=32) | |
| ap.add_argument("--micro-bs", type=int, default=8, | |
| help="GPU-memory device only: the bs-clip batch (one augmentation draw, one optimiser / EMA / scheduler step, one clipping) is pushed " | |
| "through the 42.5 M-parameter backbone in slices of this size with gradient accumulation weighted n_slice / n_batch, which equals " | |
| "the full-batch mean loss exactly (no BatchNorm in the model); 0 = full batch") | |
| ap.add_argument("--lr-backbone", type=float, default=1e-4) | |
| ap.add_argument("--lr-head", type=float, default=1e-3) | |
| ap.add_argument("--wd", type=float, default=0.01) | |
| ap.add_argument("--warmup", type=int, default=10) | |
| ap.add_argument("--ema", type=float, default=0.999) | |
| ap.add_argument("--clip", type=float, default=1.0) | |
| ap.add_argument("--loss", default="euc", choices=("euc", "l1", "mse")) | |
| ap.add_argument("--aug-window", type=int, default=1) | |
| 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("--flip-p", type=float, default=0.0) | |
| # run control (as train_kp) | |
| ap.add_argument("--max-minutes", type=float, default=60.0) | |
| ap.add_argument("--wait-free-gib", type=float, default=20.0, | |
| help="inner VRAM gate: poll torch.cuda.mem_get_info every 60 s until at least this much VRAM is free, then allocate immediately (closes " | |
| "the race between the queue's gate and the concurrent holdout extraction, which grabs ~27 GB within seconds of a chunk start); 0 = off") | |
| ap.add_argument("--yield-to", default="python.*data/holdout_pipeline", | |
| help="inner gate also waits while a process matching this regex is alive (read-only `pgrep -fc`): the holdout extraction worker ramps from " | |
| "~1 to ~27 GB over its first minutes, so free VRAM alone is not a safe launch signal; empty = off") | |
| ap.add_argument("--vram-guard-gib", type=float, default=10.0, | |
| help="VRAM turn-taking guard: a reserved tensor of this size held for the whole run so that, together with the model (~9 GiB), less than the " | |
| "holdout queue's launch threshold (14 GB) stays free -> it waits for the next inter-run gap instead of starting a 27-GB chunk mid-run " | |
| "and OOM-crashing (its own or this process); released before the CPU bootstrap; 0 = off") | |
| ap.add_argument("--n-boot", type=int, default=10000) | |
| ap.add_argument("--refs", default=None, help="name=run,... for paired Δ (default: ZV-shrink / SiMLPe / CockpitWM / CockpitWM+BATON / MotionBERT-kp of the same seed)") | |
| 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: first N train clips only") | |
| ap.add_argument("--out-root", default=None) | |
| ap.add_argument("--no-csv", action="store_true") | |
| ap.add_argument("--defer-eval", action="store_true") | |
| ap.add_argument("--eval-only", action="store_true") | |
| a = ap.parse_args(argv) | |
| a.model = MODEL_NAME[(a.formulation, a.init)] | |
| if a.refs is None: | |
| a.refs = ",".join(r.format(seed=a.seed) for r in DEFAULT_REFS) | |
| return a | |
| def procs_alive(pattern: str) -> int: | |
| """number of processes whose command line matches `pattern` (observation only; never signals anything).""" | |
| import subprocess | |
| try: | |
| out = subprocess.run(["pgrep", "-fc", pattern], capture_output=True, text=True, timeout=10) | |
| return int((out.stdout or "0").strip() or 0) | |
| except Exception: | |
| return 0 | |
| 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}_v2" | |
| 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] | |
| 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) | |
| device = torch.device("cpu" if a.cpu else "cuda") | |
| guard, waited_s = None, 0.0 | |
| if device.type == "cuda": | |
| t_w = time.time() | |
| while a.wait_free_gib > 0: # inner gate, re-checked right before the first allocation | |
| free_gib = torch.cuda.mem_get_info()[0] / 2 ** 30 | |
| busy = procs_alive(a.yield_to) if a.yield_to else 0 | |
| if free_gib >= a.wait_free_gib and busy == 0: | |
| if a.vram_guard_gib > 0: | |
| guard = torch.empty(int(a.vram_guard_gib * 2 ** 30), dtype=torch.uint8, device=device) # turn-taking guard (see --vram-guard-gib) | |
| time.sleep(5) | |
| if a.yield_to and procs_alive(a.yield_to) > 0: # the other worker launched in the same window -> release and yield | |
| guard = None | |
| torch.cuda.empty_cache() | |
| print(f"[{name}] vram gate YIELD: `{a.yield_to}` started in the same window; released the guard, sleep 60 s", flush=True) | |
| time.sleep(60) | |
| continue | |
| break | |
| print(f"[{name}] vram gate WAIT: free {free_gib:.1f} GiB (need {a.wait_free_gib} GiB), `{a.yield_to}` alive: {busy}; sleep 60 s", flush=True) | |
| time.sleep(60) | |
| waited_s = time.time() - t_w | |
| if a.wait_free_gib <= 0 and a.vram_guard_gib > 0: | |
| guard = torch.empty(int(a.vram_guard_gib * 2 ** 30), dtype=torch.uint8, device=device) | |
| print(f"[{name}] vram gate OK after {waited_s:.0f} s: free {torch.cuda.mem_get_info()[0] / 2 ** 30:.1f} GiB after reserving the {a.vram_guard_gib} GiB guard", flush=True) | |
| t_gpu = time.time() | |
| model = A.build(a.formulation, a.init, ckpt=a.ckpt) | |
| init_report = A.weight_summary(model) | |
| (out_dir / "weights_report.json").write_text(json.dumps(init_report, indent=1, default=str)) | |
| model = model.to(device) | |
| if a.init == "released": | |
| print(f"[{name}] released weights {a.ckpt}: {init_report['n_loaded_exact']}/{init_report['n_ckpt_tensors']} tensors loaded exactly " | |
| f"({init_report['n_params_loaded']:,} params), dropped {init_report['dropped']}, new {init_report['new_parameters']}, " | |
| f"copied {init_report['copied_parameters']}", flush=True) | |
| else: | |
| print(f"[{name}] RANDOM init (official _init_weights); read-out {a.formulation}", flush=True) | |
| n_params = sum(p.numel() for p in model.parameters()) | |
| ema = EMA(model, a.ema) | |
| pg = model.param_groups() | |
| opt = torch.optim.AdamW([{"params": pg["backbone"], "lr": a.lr_backbone}, {"params": pg["head"], "lr": a.lr_head}], 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) | |
| rng = np.random.default_rng(a.seed) | |
| w = {"euc": 1.0, "nll": 0.0, "bone": 0.0, "jerk": 0.0} | |
| if device.type == "cuda": | |
| torch.cuda.reset_peak_memory_stats() | |
| 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)} val ZV All-MPJPE {zv_va:.3f} " | |
| f"lr backbone {a.lr_backbone} head {a.lr_head} wd {a.wd} bs {a.bs} epochs<={a.epochs} patience {a.patience} loss {a.loss} " | |
| f"aug noise {a.noise_px} drop {a.joint_drop} tr {a.translate} sc {a.scale}", flush=True) | |
| for ep in range(a.epochs): | |
| model.train() | |
| perm = rng.permutation(len(S45_tr)) | |
| acc, nb = 0.0, 0 | |
| for i in range(0, len(perm), a.bs): | |
| obs, fut = sampler.batch(perm[i:i + a.bs], device) # ONE augmentation draw for the whole bs-clip batch (as train_kp) | |
| opt.zero_grad(set_to_none=True) | |
| n_b = obs.shape[0] | |
| mbs = n_b if a.micro_bs <= 0 else a.micro_bs | |
| loss_val = 0.0 | |
| for j in range(0, n_b, mbs): # exact gradient accumulation: sum_j (n_j / n_b) * mean_j == full-batch mean | |
| sl = slice(j, min(j + mbs, n_b)) | |
| out = model(obs[sl]) | |
| l_j, _ = LS.compute_losses(out, fut[sl], w, a.loss) | |
| (l_j * ((sl.stop - sl.start) / n_b)).backward() | |
| loss_val += l_j.item() * ((sl.stop - sl.start) / n_b) | |
| loss = torch.tensor(loss_val) | |
| if a.clip > 0: | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), a.clip) | |
| opt.step() | |
| sched.step() | |
| ema.update(model) | |
| acc += float(loss) | |
| nb += 1 | |
| pv = predict(ema.shadow, S10_va, device, bs=64) # small eval slices: keeps the VRAM footprint low next to the holdout extraction | |
| 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], "lr_head": sched.get_last_lr()[1], | |
| "t": round(time.time() - t_start, 1), "train_total": acc / nb} | |
| 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, "args": vars(a)}, out_dir / "ckpt_best.pt") | |
| else: | |
| bad += 1 | |
| if ep % 5 == 0 or improved or ep == a.epochs - 1: | |
| print(f"ep {ep:3d} loss {acc / nb:7.3f} val All {v_all:7.3f} HM {v_hm:7.2f} best {best:7.3f}@{best_ep} lr {sched.get_last_lr()[0]:.2e} " | |
| f"{(time.time() - t_start) / 60:.1f}min", flush=True) | |
| if bad >= a.patience: | |
| stop_reason = 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, weights_only=False) | |
| ema.shadow.load_state_dict(ck["model"]) | |
| preds = {} | |
| for split, S10, ids in (("val", S10_va, id_va), ("test", S10_te, id_te)): | |
| pr = predict(ema.shadow, S10, device, bs=64) | |
| 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)) | |
| preds[split] = out_dir / f"pred_{split}.npz" | |
| peak = float(torch.cuda.max_memory_allocated() / 2 ** 30) if device.type == "cuda" else None # includes the guard tensor | |
| peak_model = None if peak is None else peak - (a.vram_guard_gib if guard is not None else 0.0) | |
| if guard is not None: | |
| del guard | |
| torch.cuda.empty_cache() | |
| 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 = {} | |
| 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} | test All {quick['test']['mpjpe_all']:.3f} " | |
| f"HM {quick['test']['mpjpe_hm']:.2f}; GPU {gpu_minutes:.2f} min; peak {peak if peak is None else round(peak, 2)} GiB " | |
| f"(model only {peak_model if peak_model is None else round(peak_model, 2)} GiB, guard {a.vram_guard_gib} GiB); {stop_reason}", flush=True) | |
| code_files = [__file__, A.__file__, str(Path(A.__file__).with_name("DSTformer.py")), str(A.JOINT_MAP_JSON)] | |
| # evaluate_and_record (frozen train_kp) forwards a fixed key set into the run json -> persist the extra facts of this trainer here | |
| (out_dir / "train_info.json").write_text(json.dumps({ | |
| "name": name, "gpu_minutes": round(gpu_minutes, 2), "peak_gpu_gib": None if peak is None else round(peak, 3), | |
| "peak_model_gib": None if peak_model is None else round(peak_model, 3), "vram_guard_gib": a.vram_guard_gib, "vram_gate_wait_s": round(waited_s, 1), | |
| "micro_bs": a.micro_bs, "bs": a.bs, "best_epoch": best_ep, "epochs_run": len(hist), "stop_reason": stop_reason, "best_val_mpjpe_all": best, | |
| "quick": quick, "formulation": a.formulation, "init": a.init, "ckpt": a.ckpt}, indent=1)) | |
| return {"name": name, "out_dir": str(out_dir), "args": vars(a), "code_files": 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), "peak_gpu_gib": None if peak is None else round(peak, 3), | |
| "peak_model_gib": None if peak_model is None else round(peak_model, 3), "vram_guard_gib": a.vram_guard_gib, "vram_gate_wait_s": round(waited_s, 1), | |
| "device": str(device), "t_start": t_start, | |
| "quick": quick, "selection": "mpjpe", "model_report": None, "motion_oversample": None, "best_val_f1_mean": None, "best_f1_epoch": None, | |
| "preds": {k: str(v) for k, v in preds.items()}, "f1sel": None, "init": init_report, "hazard": None} | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 19.2 kB
- Xet hash:
- f560e75d3a2bea5778e3ee56e62dd5900175223e847e5136202d854d7710803e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.