Buckets:
| """SiMLPe on AIDE keypoints (T_OBS=5 -> T_PRED=5, J=136, residual over the last observed pose). | |
| Port of dm_bench/loader/bench/models.py::SiMLPe (siMLPe, Guo et al. WACV'23 style: | |
| spatial linear -> transposed temporal linear -> residual LayerNorm+Linear blocks -> temporal | |
| read-out). Differences from the dm_bench version, all forced by the AIDE protocol: | |
| * fixed T = 5 observed frames (no padding/cropping), H = 5, J = 136, no head-pose channel; | |
| * per-frame input features = [(xy - xy_last)/scale, x/W, y/H, conf] (5*J features); | |
| * output = last observed xy + scale * residual (pixels). | |
| Interface: forward(obs [B,5,136,3]) -> pred [B,5,136,2] px. The model only ever sees `obs`. | |
| """ | |
| 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 | |
| class SiMLPeKP(nn.Module): | |
| name = "simlpe_kp" | |
| def __init__(self, d=256, blocks=12, temp_hidden=64, scale=200.0, T=P.T_OBS, H=P.T_PRED, J=P.J, | |
| use_conf=True, dropout=0.0): | |
| super().__init__() | |
| self.T, self.H, self.J, self.scale, self.use_conf = T, H, J, scale, use_conf | |
| din = J * (5 if use_conf else 4) | |
| self.spat_in = nn.Linear(din, d) | |
| self.temp_in = nn.Linear(T, temp_hidden) | |
| self.blocks = nn.ModuleList( | |
| nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.Dropout(dropout)) for _ in range(blocks)) | |
| self.temp_out = nn.Linear(temp_hidden, H) | |
| self.out = nn.Linear(d, J * 2) | |
| nn.init.zeros_(self.out.weight) # start exactly at zero-velocity | |
| nn.init.zeros_(self.out.bias) | |
| def features(self, obs): | |
| xy = obs[..., :2] | |
| last = xy[:, -1:] # [B,1,J,2] | |
| rel = (xy - last) / self.scale | |
| ab = torch.stack([xy[..., 0] / P.W, xy[..., 1] / P.H], -1) | |
| f = [rel, ab] | |
| if self.use_conf: | |
| f.append(obs[..., 2:3]) | |
| return torch.cat(f, -1).flatten(2), last # [B,T,din] | |
| def forward(self, obs): | |
| assert obs.shape[1] == self.T and obs.shape[2] == self.J, obs.shape | |
| x, last = self.features(obs) | |
| x = self.spat_in(x) # [B,T,d] | |
| x = self.temp_in(x.transpose(1, 2)).transpose(1, 2) # [B,th,d] | |
| for blk in self.blocks: | |
| x = x + blk(x) | |
| y = self.temp_out(x.transpose(1, 2)).transpose(1, 2) # [B,H,d] | |
| res = self.out(y).view(-1, self.H, self.J, 2) | |
| return last + self.scale * res | |
| def build(name: str, **kw) -> nn.Module: | |
| if name == "simlpe_kp": | |
| return SiMLPeKP(**kw) | |
| if name == "simlpe_kp_small": | |
| return SiMLPeKP(d=128, blocks=6, **kw) | |
| raise ValueError(name) | |
| if __name__ == "__main__": | |
| m = SiMLPeKP() | |
| n = sum(p.numel() for p in m.parameters()) | |
| x = torch.rand(2, P.T_OBS, P.J, 3) * torch.tensor([P.W, P.H, 1.0]) | |
| y = m(x) | |
| print("params", n, "out", tuple(y.shape), "init == zero-velocity:", | |
| torch.allclose(y, x[:, -1:, :, :2].expand_as(y))) | |
Xet Storage Details
- Size:
- 3.19 kB
- Xet hash:
- 4751685b82fe5d90da2e34bf3eaf4b4936b0e11cb136c0232bbac6b8a156697b
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.