Buckets:
| """MotionBERT-style DSTformer on AIDE keypoints (roadmap supplement B, item D: fair baseline on the keypoint protocol). | |
| Zhu et al., MotionBERT (ICCV'23): the DSTformer alternates spatial self-attention (over the J joints of one frame) | |
| and temporal self-attention (over the T frames of one joint) with pre-norm residual MLPs. This port keeps that | |
| alternation (S -> T per block; the paper's dual-stream S/T fusion is dropped for the 5-frame protocol) and adapts | |
| the interface to the locked AIDE protocol exactly like `simlpe_kp.py`: | |
| * input obs [B,5,136,3] px + conf -> per-joint tokens [(Δx,Δy)/SCALE to the last observed pose, x/W, y/H, conf] | |
| -> Linear(5 -> d) + joint embedding + frame embedding (d = 128 by default) | |
| * 2-3 blocks of {spatial MHA over J, temporal MHA over T, MLP}, pre-norm, residual | |
| * read-out: temporal Linear(T -> H) on the per-joint tokens (SiMLPe-style), then Linear(d -> 2) per joint | |
| * output pred = last observed xy + SCALE * residual; the output layer is ZERO-initialised so the model is exactly | |
| zero-velocity at initialisation (same guarantee as CWM3-D / SiMLPe-kp; `aide/tests/test_baselines_kp.py`). | |
| Only the T_OBS observed frames are ever seen (forward asserts the shape); tensor output -> plain Euclidean loss. | |
| Formulation ablation (docs/decisions.md 2026-09-15, explanatory): `formulation="abs"` -> the read-out Linear keeps PyTorch's default | |
| init and emits ABSOLUTE normalised coordinates, pred = out(y) ⊙ (W, H) px (no residual to the last pose, no ×SCALE, no zero-velocity | |
| start); tokens / blocks / temporal read-out unchanged. `formulation="residual"` (default) is the unchanged model. | |
| """ | |
| from __future__ import annotations | |
| import sys | |
| from pathlib import Path | |
| import torch | |
| import torch.nn as nn | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) | |
| from aide.data import protocol as P # noqa: E402 | |
| SCALE = 200.0 | |
| N_CH = 5 | |
| def joint_tokens(obs: torch.Tensor) -> torch.Tensor: | |
| """obs [B,T,J,3] -> [B,T,J,5] (same channels as CWM3-D / SiMLPe-kp).""" | |
| xy = obs[..., :2] | |
| rel = (xy - xy[:, -1:]) / SCALE | |
| ab = torch.stack([xy[..., 0] / P.W, xy[..., 1] / P.H], -1) | |
| return torch.cat([rel, ab, obs[..., 2:3]], -1) | |
| class Attn(nn.Module): | |
| def __init__(self, d, heads, dropout=0.0): | |
| super().__init__() | |
| self.ln = nn.LayerNorm(d) | |
| self.attn = nn.MultiheadAttention(d, heads, dropout=dropout, batch_first=True) | |
| def forward(self, x): # x [N, L, d] -> residual attention over L | |
| h = self.ln(x) | |
| return x + self.attn(h, h, h, need_weights=False)[0] | |
| class DSTBlock(nn.Module): | |
| """Spatial attention (over joints, per frame) -> temporal attention (over frames, per joint) -> MLP.""" | |
| def __init__(self, d, heads, ffn, dropout=0.0): | |
| super().__init__() | |
| self.spatial = Attn(d, heads, dropout) | |
| self.temporal = Attn(d, heads, dropout) | |
| self.ln = nn.LayerNorm(d) | |
| self.ff = nn.Sequential(nn.Linear(d, ffn), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn, d)) | |
| def forward(self, x): # x [B,T,J,d] | |
| B, T, J, d = x.shape | |
| x = self.spatial(x.reshape(B * T, J, d)).view(B, T, J, d) | |
| x = self.temporal(x.permute(0, 2, 1, 3).reshape(B * J, T, d)).view(B, J, T, d).permute(0, 2, 1, 3) | |
| return x + self.ff(self.ln(x)) | |
| class MotionBERTKP(nn.Module): | |
| name = "motionbert_kp" | |
| # trunk-only initialisation (arch_attribution (1), train_kp --init-parts trunk): input projection / joint + frame embeddings / | |
| # DST blocks / final LN are the trunk; the read-out head (temp_out, out) keeps its seed init (out is zero -> the fine-tune still | |
| # starts exactly at zero-velocity), mirroring cwm3d's tokenizer / blocks / ln_f vs horizon-head split. | |
| TRUNK_PREFIXES = ("inp.", "joint_emb", "time_emb", "blocks.", "ln_f.") | |
| def __init__(self, d=128, blocks=3, heads=4, ffn=None, dropout=0.0, T=P.T_OBS, H=P.T_PRED, J=P.J, formulation="residual"): | |
| super().__init__() | |
| assert formulation in ("residual", "abs"), formulation | |
| ffn = 4 * d if ffn is None else ffn | |
| self.T, self.H, self.J, self.formulation = T, H, J, formulation | |
| self.inp = nn.Linear(N_CH, d) | |
| self.joint_emb = nn.Parameter(torch.zeros(J, d)) | |
| self.time_emb = nn.Parameter(torch.zeros(T, d)) | |
| nn.init.normal_(self.joint_emb, std=0.02) | |
| nn.init.normal_(self.time_emb, std=0.02) | |
| self.blocks = nn.ModuleList(DSTBlock(d, heads, ffn, dropout) for _ in range(blocks)) | |
| self.ln_f = nn.LayerNorm(d) | |
| self.temp_out = nn.Linear(T, H) # temporal read-out (observed frames -> horizons) | |
| self.out = nn.Linear(d, 2) | |
| if formulation == "residual": | |
| nn.init.zeros_(self.out.weight) # exactly zero-velocity at init | |
| nn.init.zeros_(self.out.bias) | |
| else: # abs: default init kept; (W, H) buffer non-persistent -> same state dict keys | |
| self.register_buffer("wh", torch.tensor([float(P.W), float(P.H)]), persistent=False) | |
| def forward(self, obs): | |
| assert obs.shape[1] == self.T and obs.shape[2] == self.J, obs.shape | |
| last = obs[:, -1:, :, :2] | |
| x = self.inp(joint_tokens(obs)) + self.joint_emb[None, None] + self.time_emb[None, :, None] # [B,T,J,d] | |
| for blk in self.blocks: | |
| x = blk(x) | |
| x = self.ln_f(x) | |
| y = self.temp_out(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) # [B,J,d,T] -> [B,J,d,H] -> [B,H,J,d] | |
| if self.formulation == "abs": | |
| return self.out(y) * self.wh # absolute normalised coordinates -> px | |
| return last + SCALE * self.out(y) | |
| def parse_variant(name: str) -> dict: | |
| """Name-encoded width / depth variants for the K-v2 programme's MB fairness grid (protocol C23, S3 i): | |
| `motionbert_kp[_d<width>][_L<blocks>]`, e.g. motionbert_kp_d192, motionbert_kp_L4, motionbert_kp_d192_L4 -> {"d": 192, "blocks": 4}. | |
| The plain name and `motionbert_kp_2b` are unchanged; any other suffix is rejected (run-name tags such as lr / wd / do are trainer flags, not model names).""" | |
| if name == "motionbert_kp": | |
| return {} | |
| if name == "motionbert_kp_2b": | |
| return {"blocks": 2} | |
| if not name.startswith("motionbert_kp_"): | |
| raise ValueError(name) | |
| kw = {} | |
| for tok in name[len("motionbert_kp_"):].split("_"): | |
| if tok.startswith("d") and tok[1:].isdigit() and "d" not in kw: | |
| kw["d"] = int(tok[1:]) | |
| elif tok.startswith("L") and tok[1:].isdigit() and "blocks" not in kw: | |
| kw["blocks"] = int(tok[1:]) | |
| else: | |
| raise ValueError(f"unknown motionbert_kp variant token {tok!r} in {name!r}") | |
| return kw | |
| def build(name: str, **kw) -> nn.Module: | |
| variant = parse_variant(name) # ValueError for unknown names (unchanged behaviour for the two original names) | |
| clash = set(variant) & set(kw) | |
| if clash: | |
| raise ValueError(f"{name}: {sorted(clash)} given both in the name and as keyword") | |
| return MotionBERTKP(**variant, **kw) | |
| if __name__ == "__main__": | |
| m = MotionBERTKP() | |
| x = torch.rand(2, P.T_OBS, P.J, 3) * torch.tensor([P.W, P.H, 1.0]) | |
| y = m(x) | |
| print("params", sum(p.numel() for p in m.parameters()), "out", tuple(y.shape), | |
| "init == zero-velocity:", torch.equal(y, x[:, -1:, :, :2].expand_as(y))) | |
Xet Storage Details
- Size:
- 7.51 kB
- Xet hash:
- d5038db0fcac68c2b19d36d2c3d8c794d0d675042bc7d38e23ed957cedb5d930
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.