Buckets:
| """Guards for the faithful MotionBERT baseline (aide/models/motionbert_official; pre-registration docs/decisions.md 2026-09-16T21:28:48Z). | |
| * released weights loaded: the probe block tensor inside the model is bit-identical to the checkpoint (sha256), 240 block tensors + temporal | |
| position embedding + joints_embed / norm / pre_logits / ts_attn loaded, only the 3-D head dropped, only the 2-D read-out new; | |
| * (a) residual read-out is exactly zero-velocity at initialisation (released AND random init); (b) abs is not; | |
| * the model sees only the T_OBS observed frames: the 5 placeholder frames are copies of the last observed frame, perturbing the FUTURE ground | |
| truth of a 10-frame block leaves the output bit-identical, 10-frame inputs are rejected; | |
| * output shape [B, 5, 136, 2]; normalisation is the repo's --pixel map and denormalise inverts it; joint map = repo halpe2h36m for the body. | |
| python -m pytest tests/test_motionbert_official.py -q | |
| """ | |
| from __future__ import annotations | |
| import os | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| os.environ.setdefault("MKL_NUM_THREADS", "4") | |
| import hashlib # noqa: E402 | |
| import json # noqa: E402 | |
| 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[1] | |
| sys.path.insert(0, str(ROOT)) | |
| from aide.data import protocol as P # noqa: E402 | |
| from aide.models.motionbert_official import adapter as A # noqa: E402 | |
| from aide.train.train_kp import split_obs_fut # noqa: E402 | |
| CKPT = Path(A.CKPTS["MB_release"]["path"]) | |
| needs_ckpt = pytest.mark.skipif(not CKPT.exists(), reason=f"released checkpoint not downloaded: {CKPT}") | |
| def _frames(B=2, seed=0): | |
| g = torch.Generator().manual_seed(seed) | |
| return torch.rand(B, P.N_SAMPLE, P.J, 3, generator=g) * torch.tensor([P.W, P.H, 1.0]) | |
| def released(): | |
| torch.manual_seed(0) | |
| return A.build("residual", "released").eval() | |
| def test_released_weights_loaded_bit_exact(released): | |
| rep = released.load_report | |
| assert rep["sha256"] == A.CKPTS["MB_release"]["sha256"] and rep["bytes"] == A.CKPTS["MB_release"]["bytes"] | |
| assert rep["n_ckpt_tensors"] == 260 and rep["n_loaded_exact"] == 258 and rep["dropped"] == ["head.bias", "head.weight"] | |
| assert rep["loaded_groups"]["blocks_st+blocks_ts"] == 240 and rep["loaded_groups"]["temp_embed"] == 1 | |
| assert not rep["shape_mismatch"] and not rep["missing_in_ckpt"] | |
| # checksum of a loaded block tensor inside the MODEL equals the checkpoint's tensor | |
| sd, _ = A.read_checkpoint("MB_release") | |
| ck_t = sd["blocks_st.0.attn_s.qkv.weight"] | |
| model_t = released.backbone.blocks_st[0].attn_s.qkv.weight | |
| assert hashlib.sha256(model_t.detach().numpy().tobytes()).hexdigest() == hashlib.sha256(ck_t.numpy().tobytes()).hexdigest() == rep["probe_sha256"] | |
| assert torch.equal(released.backbone.temp_embed, sd["temp_embed"]) | |
| for k in ("blocks_ts.4.mlp_t.fc2.weight", "pre_logits.fc.weight", "norm.weight", "joints_embed.weight", "ts_attn.3.weight"): | |
| assert torch.equal(released.backbone.state_dict()[k], sd[k]), k | |
| # spatial position embedding: 136 rows, every row a copy of the pretrained row of its mapped H36M joint | |
| jm = A.load_joint_map() | |
| assert tuple(released.backbone.pos_embed.shape) == (1, P.J, 512) | |
| assert torch.equal(released.backbone.pos_embed[0], sd["pos_embed"][0][torch.as_tensor(jm)]) | |
| # only the read-out is new | |
| assert released.new_parameters == ["head.weight", "head.bias"] | |
| assert sum(p.numel() for p in released.parameters()) == rep["n_params_loaded"] - 17 * 512 + P.J * 512 + 512 * 2 + 2 | |
| def test_zero_velocity_identity_at_init_released(released): | |
| x = _frames() | |
| obs, _ = split_obs_fut(x) | |
| with torch.no_grad(): | |
| y = released(obs) | |
| assert y.shape == (2, P.T_PRED, P.J, 2) | |
| assert torch.equal(y, obs[:, -1:, :, :2].expand_as(y)) | |
| def test_zero_velocity_identity_at_init_random_and_abs_is_not(): | |
| torch.manual_seed(1) | |
| m = A.build("residual", "random").eval() | |
| x = _frames(seed=3) | |
| obs, _ = split_obs_fut(x) | |
| with torch.no_grad(): | |
| y = m(obs) | |
| assert torch.equal(y, obs[:, -1:, :, :2].expand_as(y)) | |
| torch.manual_seed(1) | |
| mb = A.build("abs", "random").eval() | |
| with torch.no_grad(): | |
| yb = mb(obs) | |
| assert yb.shape == (2, P.T_PRED, P.J, 2) and not torch.equal(yb, obs[:, -1:, :, :2].expand_as(yb)) | |
| def test_no_future_leakage_and_placeholders_are_last_frame(released): | |
| torch.manual_seed(2) | |
| m = A.build("abs", "released").eval() # abs read-out moves at init -> a non-trivial output to compare | |
| x = _frames(seed=5) | |
| obs, fut = split_obs_fut(x) | |
| x2 = x.clone() | |
| x2[:, P.T_OBS:] += 250.0 # perturb the FUTURE ground truth only | |
| obs2, fut2 = split_obs_fut(x2) | |
| assert torch.equal(obs, obs2) and not torch.equal(fut, fut2) | |
| with torch.no_grad(): | |
| y1, y2 = m(obs), m(obs2) | |
| inp = m.build_input(obs) | |
| assert torch.equal(y1, y2) | |
| assert inp.shape == (2, P.N_SAMPLE, P.J, 3) | |
| assert torch.equal(inp[:, P.T_OBS:], inp[:, P.T_OBS - 1:P.T_OBS].expand(-1, P.T_PRED, -1, -1)) # placeholders = last observed frame | |
| # perturbing an OBSERVED frame changes the output (the model is live) | |
| obs3 = obs.clone() | |
| obs3[:, 0, :, :2] += 50.0 | |
| with torch.no_grad(): | |
| y3 = m(obs3) | |
| assert not torch.equal(y1, y3) | |
| with pytest.raises(AssertionError): | |
| m(x) # 10-frame input rejected | |
| def test_pixel_normalisation_and_inverse_and_joint_map(): | |
| x = _frames(seed=7) | |
| n = A.normalise(x) | |
| assert torch.allclose(n[..., 1], (x[..., 1] - P.H / 2) / (min(P.W, P.H) / 2)) | |
| assert torch.allclose(n[..., 0], (x[..., 0] - P.W / 2) / (min(P.W, P.H) / 2)) | |
| assert torch.equal(n[..., 2], x[..., 2]) | |
| assert torch.allclose(A.denormalise_xy(n[..., :2]), x[..., :2], atol=1e-3) | |
| # repo halpe2h36m: H36M row m holds HALPE joint H36M_FROM_HALPE[m]; belly = midpoint neck / hip | |
| h = A.halpe2h36m_np(x.numpy()) | |
| for m_, hj in A.H36M_FROM_HALPE.items(): | |
| assert np.array_equal(h[..., m_, :], x.numpy()[..., hj, :]) | |
| assert np.allclose(h[..., 7, :], 0.5 * (x.numpy()[..., 18, :] + x.numpy()[..., 19, :])) | |
| jm = A.load_joint_map() | |
| d = json.loads(A.JOINT_MAP_JSON.read_text()) | |
| assert jm.shape == (P.J,) and len(d["rows"]) == P.J | |
| for hj, m_ in A.HALPE2H36M_DIRECT.items(): | |
| assert jm[hj] == m_ # body joints via halpe2h36m | |
| assert (jm[94:115] == 13).all() and (jm[115:136] == 16).all() # hands <- wrists | |
| assert (jm[[20, 22, 24]] == 6).all() and (jm[[21, 23, 25]] == 3).all() # feet <- ankles | |
| assert set(jm[26:94].tolist()) <= {9, 10} and set(jm[1:5].tolist()) <= {9, 10} # face / eyes / ears <- nose or head | |
Xet Storage Details
- Size:
- 6.98 kB
- Xet hash:
- 76c783a976e9b79a5868ac81a554a49a997bf83dc383176aedc3d33443aa2afc
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.