Buckets:
| """F7 (docs/decisions.md 2026-09-17 pre-registration) — AIDE-only multi-stride pretraining + ensemble distillation: additive-code guards. | |
| 1. Defaults unchanged: train_kp.get_args without the new flags -> no teacher (load_teachers returns (None, None)); LS.distillation with every | |
| weight 0 returns an exact zero and no terms; state dict / forward of the default cwm3d equal the pre-change fixture (byte identity). | |
| 2. Distillation terms: student == teacher -> every term 0; KL >= 0; gradient reaches the student only; tensor-output models get the | |
| coordinate term only; mismatched weights / teachers are rejected. | |
| 3. Teacher loading: strict state-dict load of a real adopted run reproduces its saved pred_val.npz (skipped if the run dir is absent). | |
| 4. Multi-stride grids: every index in [0, 44], uniform strides have equal gaps, the mixed family has gaps in {4, 5}, starts cover the range; | |
| isolation gate rejects BATON corpus / route / hold-out tokens and accepts the F7 argv. | |
| 5. CPU smoke of the pretraining script (8 clips, 2 epochs) writes ckpt_best.pt / history.json / TRAIN_DONE and the checkpoint loads into | |
| train_kp's --init-parts trunk path (59 cwm3d trunk tensors, horizon head kept). | |
| python -m pytest aide/tests/test_f7_distil_pretrain.py -q | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import os | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import sys # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import pytest # noqa: E402 | |
| import torch # noqa: E402 | |
| torch.set_num_threads(4) | |
| ROOT = Path(__file__).resolve().parents[2] | |
| sys.path.insert(0, str(ROOT)) | |
| from aide.data import protocol as P # noqa: E402 | |
| from aide.models import losses as LS # noqa: E402 | |
| from aide.train import pretrain_aide_multistride as PA # noqa: E402 | |
| from aide.train.train_kp import build_model, get_args, load_init_ckpt, load_teachers, model_pred, predict, split_obs_fut, teacher_target # noqa: E402 | |
| REF = ROOT / "aide" / "tests" / "fixtures" / "formulation_abs_prechange.pt" | |
| RUNS = P.RESULTS / "runs" | |
| ADOPTED = "cwm3d_kp_batonpre_trunkp123_s42_v2" | |
| def _sha(m): | |
| h = hashlib.sha256() | |
| for k, v in sorted(m.state_dict().items()): | |
| h.update(k.encode()); h.update(str(tuple(v.shape)).encode()); h.update(v.detach().cpu().contiguous().numpy().tobytes()) | |
| return h.hexdigest() | |
| def _frames(n=4, seed=0): | |
| g = torch.Generator().manual_seed(seed) | |
| return torch.rand(n, P.N_SAMPLE, P.J, 3, generator=g) * torch.tensor([P.W, P.H, 1.0]) | |
| def _model(seed=0, name="cwm3d", extra=()): | |
| a = get_args(["--model", name, "--cpu", *extra]) | |
| torch.manual_seed(seed) | |
| m, _ = build_model(name, a) | |
| return m, a | |
| # ----------------------------------------------------------------------------- 1. defaults | |
| def test_defaults_no_teacher_and_zero_distil(): | |
| m, a = _model() | |
| assert a.distil_teachers == "" and a.w_distil == 0.0 and a.w_distil_modes == 0.0 and a.w_distil_logb == 0.0 | |
| assert load_teachers(a, torch.device("cpu")) == (None, None) | |
| out = m(split_obs_fut(_frames())[0]) | |
| z, terms = LS.distillation(out, {"pred": out["pred"]}, 0.0, 0.0, 0.0) | |
| assert float(z) == 0.0 and terms == {} | |
| z2, t2 = LS.distillation(out["pred"], {"pred": out["pred"]}, 0.0, 0.0, 0.0) | |
| assert float(z2) == 0.0 and t2 == {} | |
| def test_default_cwm3d_byte_identical_to_prechange_fixture(): | |
| ref = torch.load(REF, map_location="cpu", weights_only=False) | |
| m, _ = _model(seed=0) | |
| assert sorted(m.state_dict()) == sorted(ref["state_dict"]) if "state_dict" in ref else True | |
| if "state_dict" in ref: | |
| for k, v in ref["state_dict"].items(): | |
| assert torch.equal(m.state_dict()[k], v), k | |
| if "sha256" in ref: | |
| assert _sha(m) == ref["sha256"] | |
| # ----------------------------------------------------------------------------- 2. distillation terms | |
| def test_distillation_zero_at_equality_and_positive_otherwise(): | |
| m, _ = _model(seed=1) | |
| with torch.no_grad(): | |
| for p in m.parameters(): | |
| p.add_(0.02 * torch.randn(p.shape, generator=torch.Generator().manual_seed(2))) | |
| obs = split_obs_fut(_frames())[0] | |
| out = m(obs) | |
| tgt = teacher_target([m], obs) | |
| l, terms = LS.distillation(out, tgt, 1.0, 0.1, 0.1) | |
| assert set(terms) == {"distil", "distil_modes", "distil_logb"} | |
| assert abs(float(l.detach())) < 2e-3 and all(abs(float(v)) < 2e-3 for v in terms.values()) # euclidean eps floor sqrt(1e-6) = 1e-3 px | |
| tgt2 = {"pred": tgt["pred"] + 5.0, "log_b": tgt["log_b"] + 0.5, "anchor_w": torch.softmax(torch.randn_like(tgt["anchor_w"]), -1)} | |
| l2, t2 = LS.distillation(out, tgt2, 1.0, 0.1, 0.1) | |
| assert abs(float(t2["distil"]) - 5.0 * 2 ** 0.5) < 1e-3 # +5 px on x and y -> 5*sqrt(2) and float(t2["distil_modes"]) > 0 and abs(float(t2["distil_logb"]) - 0.5) < 1e-4 | |
| assert float(l2) > 0 | |
| def test_distillation_gradient_reaches_student_only(): | |
| s, _ = _model(seed=3) | |
| t, _ = _model(seed=4) | |
| with torch.no_grad(): # both models are zero-velocity at init; move the teacher so the target differs | |
| for p in t.parameters(): | |
| p.add_(0.05 * torch.randn(p.shape, generator=torch.Generator().manual_seed(9))) | |
| for p in t.parameters(): | |
| p.requires_grad_(False) | |
| obs = split_obs_fut(_frames())[0] | |
| out = s(obs) | |
| tgt = teacher_target([t, t], obs) | |
| assert not tgt["pred"].requires_grad | |
| l, _ = LS.distillation(out, tgt, 1.0, 0.1, 0.1) | |
| l.backward() | |
| assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in s.parameters()) | |
| assert all(p.grad is None for p in t.parameters()) | |
| def test_tensor_output_student_gets_coordinate_term_only(): | |
| m, _ = _model(seed=5, name="motionbert_kp") | |
| obs = split_obs_fut(_frames())[0] | |
| out = m(obs) | |
| assert torch.is_tensor(out) | |
| tgt = teacher_target([m], obs) | |
| assert tgt["log_b"] is None and tgt["anchor_w"] is None | |
| l, terms = LS.distillation(out, {"pred": tgt["pred"] + 2.0, "log_b": None, "anchor_w": None}, 1.0, 0.1, 0.1) | |
| assert set(terms) == {"distil"} and abs(float(l.detach()) - 2.0 * 2 ** 0.5) < 1e-3 | |
| def test_teacher_mean_is_the_mean_of_members(): | |
| a, _ = _model(seed=6) | |
| b, _ = _model(seed=7) | |
| with torch.no_grad(): | |
| for p in b.parameters(): | |
| p.add_(0.05 * torch.randn(p.shape, generator=torch.Generator().manual_seed(8))) | |
| obs = split_obs_fut(_frames())[0] | |
| tgt = teacher_target([a, b], obs) | |
| with torch.no_grad(): | |
| ref = (model_pred(a(obs)) + model_pred(b(obs))) / 2 | |
| assert torch.allclose(tgt["pred"], ref, atol=1e-5) | |
| def test_mismatched_flags_rejected(): | |
| a = get_args(["--model", "cwm3d", "--cpu", "--distil-teachers", "x"]) | |
| with pytest.raises(AssertionError): | |
| load_teachers(a, torch.device("cpu")) | |
| a = get_args(["--model", "cwm3d", "--cpu", "--w-distil", "1"]) | |
| with pytest.raises(AssertionError): | |
| load_teachers(a, torch.device("cpu")) | |
| # ----------------------------------------------------------------------------- 3. real teacher | |
| def test_loaded_teacher_reproduces_saved_val_predictions(): | |
| a = get_args(["--model", "cwm3d", "--cpu", "--distil-teachers", ADOPTED, "--w-distil", "1"]) | |
| teachers, rep = load_teachers(a, torch.device("cpu")) | |
| assert len(teachers) == 1 and rep["n_teachers"] == 1 and rep["teachers"][0]["run"] == ADOPTED | |
| c = np.load(P.CACHE_NPZ, allow_pickle=False) | |
| va = c["split"] == "val" | |
| S10 = c["S10"][va][:64] | |
| pv = predict(teachers[0], S10, torch.device("cpu")) | |
| saved = np.load(RUNS / ADOPTED / "pred_val.npz", allow_pickle=False) | |
| ids_saved = saved["clip_id"].astype(str) | |
| order = [int(np.nonzero(ids_saved == i)[0][0]) for i in c["clip_id"][va][:64].astype(str)] | |
| assert np.abs(pv - saved["pred"][order]).max() < 5e-3 | |
| # ----------------------------------------------------------------------------- 4. grids + isolation | |
| def test_multistride_grids_are_valid(): | |
| rng = np.random.default_rng(0) | |
| strides = PA.parse_strides("1,2,3,4,mixed") | |
| g = PA.multistride_grids(rng, 4000, strides) | |
| assert g.shape == (4000, P.N_SAMPLE) and g.min() >= 0 and g.max() <= P.N_FRAMES_STD - 1 | |
| gaps = np.diff(g, axis=1) | |
| assert (gaps > 0).all() | |
| uniform = (gaps == gaps[:, :1]).all(1) | |
| assert set(np.unique(gaps[uniform, 0]).tolist()) == {1, 2, 3, 4} | |
| assert set(np.unique(gaps[~uniform]).tolist()) <= {4, 5} | |
| for st in (1, 2, 3, 4): | |
| rows = uniform & (gaps[:, 0] == st) | |
| starts = g[rows, 0] | |
| assert starts.min() == 0 and starts.max() == P.N_FRAMES_STD - 1 - st * (P.N_SAMPLE - 1) | |
| with pytest.raises(AssertionError): | |
| PA.parse_strides("5") | |
| def test_isolation_gate(): | |
| assert PA.assert_aide_only(["--seed", "123", "--name", "cwm3d_kp_aidepre_pretrain_s123", "--init-ckpt", | |
| "aide/results/runs/cwm3d_kp_batonpre_pretrain_s123/ckpt_best.pt"]) | |
| for bad in ("aide/results/cache/baton_pose_corpus.npz", "results/pose_causal_routes_v2/x.npz", "results/holdout/pack", "ACURA/route_0-7"): | |
| with pytest.raises(AssertionError): | |
| PA.assert_aide_only(["--x", bad]) | |
| ho = ROOT / "results" / "holdout" / "holdout_leaves.txt" | |
| if ho.exists(): | |
| dev = ho.read_text().splitlines()[0].split("/")[1] | |
| with pytest.raises(AssertionError): | |
| PA.assert_aide_only(["--x", f"/some/{dev}/thing"]) | |
| # ----------------------------------------------------------------------------- 5. smoke | |
| def test_pretrain_smoke_and_trunk_init(tmp_path): | |
| PA.main(["--seed", "0", "--name", "smoke_aidepre", "--out-root", str(tmp_path), "--epochs", "2", "--windows-per-clip", "2", "--bs", "8", | |
| "--limit-train", "8", "--cpu", "--max-minutes", "5"]) | |
| d = tmp_path / "smoke_aidepre" | |
| assert (d / "ckpt_best.pt").exists() and (d / "history.json").exists() and (d / "TRAIN_DONE").exists() | |
| hist = json.loads((d / "history.json").read_text()) | |
| assert len(hist) == 2 and all(np.isfinite(h["val_mpjpe_all"]) for h in hist) | |
| rj = sorted(d.glob("run_smoke_aidepre_2*.json")) | |
| assert rj and json.load(open(rj[-1]))["data"]["split"] == "train" | |
| m, _ = _model(seed=42) | |
| rep = load_init_ckpt(m, d / "ckpt_best.pt", "trunk", "cpu") | |
| assert rep["n_loaded"] == 59 and rep["n_skipped_part"] > 0 | |
| obs = split_obs_fut(_frames())[0] | |
| with torch.no_grad(): | |
| out = m(obs) | |
| assert torch.allclose(out["pred"], obs[:, -1:, :, :2].expand(-1, P.T_PRED, -1, -1), atol=1e-4) # head at zero init -> zero-velocity start | |
Xet Storage Details
- Size:
- 10.7 kB
- Xet hash:
- 3688ed6652627d20193702ec6df351d5b3bc50decbb8472efc196d83e3632c0c
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.