Buckets:
| """ST-GCN on AIDE keypoints (roadmap supplement B, item D: fair baseline on the keypoint protocol). | |
| Reuses `STGCNBlock` / `normalized_adjacency` from `aide/models/driverwm_repro.py` (spatial graph convolution over the | |
| normalised HALPE-136 adjacency `protocol.EDGES_ALL` + self loops, temporal conv k = 3, residual, LayerNorm) so the | |
| block is literally the one inside our Driver-WM reproduction's skeleton decoder. Interface as `simlpe_kp.py`: | |
| * input obs [B,5,136,3] -> per-joint channels [(Δx,Δy)/SCALE to the last observed pose, x/W, y/H, conf] | |
| -> Linear(5 -> c) per joint (c = 128), frame embedding | |
| * 8 ST-GCN blocks over [B, T, J, c] (temporal conv along the 5 observed frames, zero-padded) | |
| * read-out: temporal Linear(T -> H), Linear(c -> 2) per joint, ZERO-initialised -> exactly zero-velocity at init | |
| * output pred = last observed xy + SCALE * residual (px) | |
| Only the T_OBS observed frames are ever seen; tensor output -> plain Euclidean loss. | |
| """ | |
| 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 | |
| from aide.models.driverwm_repro import STGCNBlock, normalized_adjacency # noqa: E402 | |
| from aide.models.motionbert_kp import N_CH, SCALE, joint_tokens # noqa: E402 | |
| class STGCNKP(nn.Module): | |
| name = "stgcn_kp" | |
| # trunk-only initialisation (arch_attribution (1), train_kp --init-parts trunk): input projection / frame embedding / ST-GCN blocks | |
| # are the trunk; the read-out head (temp_out, out) keeps its seed init (out is zero -> exact zero-velocity start of the fine-tune). | |
| TRUNK_PREFIXES = ("inp.", "time_emb", "blocks.") | |
| def __init__(self, c=128, blocks=8, T=P.T_OBS, H=P.T_PRED, J=P.J, edges=P.EDGES_ALL, dropout=0.0): | |
| super().__init__() | |
| self.T, self.H, self.J = T, H, J | |
| self.inp = nn.Linear(N_CH, c) | |
| self.time_emb = nn.Parameter(torch.zeros(T, c)) | |
| nn.init.normal_(self.time_emb, std=0.02) | |
| A = normalized_adjacency(J, edges) | |
| self.blocks = nn.ModuleList(STGCNBlock(c, A) for _ in range(blocks)) | |
| self.drop = nn.Dropout(dropout) | |
| self.temp_out = nn.Linear(T, H) | |
| self.out = nn.Linear(c, 2) | |
| nn.init.zeros_(self.out.weight) | |
| nn.init.zeros_(self.out.bias) | |
| 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.time_emb[None, :, None] # [B,T,J,c] | |
| for blk in self.blocks: | |
| x = self.drop(blk(x)) | |
| y = self.temp_out(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) # [B,H,J,c] | |
| return last + SCALE * self.out(y) | |
| def build(name: str, **kw) -> nn.Module: | |
| if name == "stgcn_kp": | |
| return STGCNKP(**kw) | |
| raise ValueError(name) | |
| if __name__ == "__main__": | |
| m = STGCNKP() | |
| 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:
- 3.22 kB
- Xet hash:
- f0d1d9da53df3f97c3df18daa3de1a4d02f46de003b3a948186c303cde36cb55
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.