Buckets:
| """CWM3-D: part-tokenised causal temporal-part Transformer with anchor-mixture horizon head | |
| (plan_2026-09-10.md §4). Keypoint-only by default; visual tokens, traffic->driver gated | |
| cross-attention and semantic heads are behind flags (default OFF) and are shape-tested only until | |
| the Qwen3-VL features exist. | |
| forward(obs [B,5,136,3], vis=None) -> {"pred":[B,5,136,2], "log_b":[B,5,136], "last":[B,1,136,2], | |
| "anchor_w":[B,5,136,3] (if anchors), "logits_<task>" (if semantic)} | |
| Rollout modes | |
| chunked (default) all 5 horizons from the last-frame part tokens in one pass | |
| recursive 5 passes; each step predicts horizon 1 only and appends it (conf = last obs conf) to a | |
| sliding 5-frame window -- the "recursion vs chunked" ablation. | |
| Zero-init guarantee: pred == last observed pose at initialisation (test_cwm3d.py). | |
| Formulation ablation (docs/decisions.md 2026-09-15, explanatory): `formulation="abs"` -> the horizon head keeps PyTorch's default init | |
| and the read-out is heads.AbsoluteReadout (pred = μ ⊙ (W, H) px: no residual to the last pose, no anchors, no zero-velocity start); | |
| `formulation="residual"` (default) is the unchanged model (test_formulation_abs.py asserts byte-identity to the pre-change code). | |
| HM h=5 arms (docs/decisions.md 2026-09-17T07:00:22Z, part (A), arms b / c): `horizon_gain="learned"|"parts"` -> heads.HorizonGain multiplies the | |
| anchored residual per horizon (g = 1 + δ, δ zero-init; "parts" = separate gains for the 42 hand joints); default None = unchanged model | |
| (no extra parameter, byte-identical; aide/tests/test_hm_h5_arms.py). | |
| """ | |
| 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.heads import (SCALE, AbsoluteReadout, AnchorMixture, Block, HorizonGain, HorizonHead, PartTokenizer, anchors, # noqa: E402 | |
| causal_part_mask) | |
| VIS_DIM = 2048 | |
| N_VIEWS = len(P.VIEWS) | |
| EXT_VIEWS = [P.VIEWS.index(v) for v in ("front", "left", "right")] | |
| EGO_DIM = 3 * 16 # 16 Farneback flow statistics x 3 exterior views (aide/features/extract_egomotion.py) | |
| class TrafficGate(nn.Module): | |
| """Gated cross-attention from driver part tokens (queries) to external-view visual tokens.""" | |
| def __init__(self, d, heads): | |
| super().__init__() | |
| self.ln_q = nn.LayerNorm(d) | |
| self.ln_kv = nn.LayerNorm(d) | |
| self.attn = nn.MultiheadAttention(d, heads, batch_first=True) | |
| self.gate = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d)) | |
| nn.init.zeros_(self.gate[-1].weight) | |
| nn.init.constant_(self.gate[-1].bias, -4.0) # sigmoid(-4) ~ 0.018: near-identity at init | |
| def forward(self, q, kv, kv_mask, override=None): | |
| """override: clamp the gate to a constant (Driver-WM Table-3 do(lambda_CA = 0 / 1 / 2)); None = learned gate.""" | |
| a = self.attn(self.ln_q(q), self.ln_kv(kv), self.ln_kv(kv), attn_mask=kv_mask, need_weights=False)[0] | |
| g = torch.sigmoid(self.gate(q)) if override is None else torch.full_like(a, float(override)) | |
| return q + g * a | |
| class SemAttnHeads(nn.Module): | |
| """Semantic heads with the Driver-WM-reproduction recipe on top of the CWM3-D trunk: per stream one | |
| pre-norm causal TransformerEncoderLayer(2048, 4 heads, d_ff 4096) over the observed-window features | |
| (in-cabin; exterior = view-embedded mean of the 3 outside views), last state -> Linear. DBR/DER also | |
| receive the driver part-token mean (d); TCR/VCR the window-mean ego-motion statistics when enabled.""" | |
| def __init__(self, d, use_ego=False, heads=4, d_ff=4096, view_gate=False, frame_avg_tasks=(), task_weights=None): | |
| super().__init__() | |
| self.d, self.view_gate, self.frame_avg_tasks = d, view_gate, tuple(frame_avg_tasks) | |
| self.task_weights = dict(task_weights or {}) | |
| if view_gate: # per-head softmax gate over the 4 view-specific causal encodings (learned, input-independent) | |
| self.gate_logits = nn.ParameterDict({t: nn.Parameter(torch.zeros(N_VIEWS)) for t in P.TASKS}) | |
| self.view_emb = nn.Parameter(torch.zeros(N_VIEWS, VIS_DIM)) | |
| self.pos = nn.Parameter(torch.zeros(P.T_OBS, VIS_DIM)) | |
| nn.init.normal_(self.view_emb, std=0.02) | |
| nn.init.normal_(self.pos, std=0.02) | |
| self.enc_in = nn.TransformerEncoderLayer(VIS_DIM, heads, d_ff, dropout=0.0, batch_first=True, norm_first=True) | |
| self.enc_ext = nn.TransformerEncoderLayer(VIS_DIM, heads, d_ff, dropout=0.0, batch_first=True, norm_first=True) | |
| self.ln_in, self.ln_ext = nn.LayerNorm(VIS_DIM), nn.LayerNorm(VIS_DIM) | |
| self.heads = nn.ModuleDict({t: nn.Linear(VIS_DIM + (d if t in ("dbr", "der") else (EGO_DIM if use_ego else 0)), | |
| len(P.CLASSES[t])) for t in P.TASKS}) | |
| def view_weights(self): | |
| return {t: torch.softmax(g.detach(), 0).tolist() for t, g in self.gate_logits.items()} if self.view_gate else None | |
| def forward(self, vis, drv, ego_stats=None): # vis [B,T,4,2048], drv [B,d] (or None), ego_stats [B,48] or None | |
| B, T = vis.shape[:2] | |
| mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=vis.device), 1) | |
| drv_cat = (lambda h: h) if drv is None else (lambda h: torch.cat([h, drv], -1)) | |
| out = {} | |
| if self.view_gate: | |
| # every view through the in-cabin encoder (shared weights + view embedding) -> per-view last states | |
| z = (vis + self.view_emb[None, None] + self.pos[None, :T, None]).permute(0, 2, 1, 3).reshape(B * N_VIEWS, T, VIS_DIM) | |
| hv_seq = self.enc_in(z, src_mask=mask).view(B, N_VIEWS, T, VIS_DIM) | |
| hv = self.ln_in(hv_seq[:, :, -1]) # [B,4,2048] | |
| for t in P.TASKS: | |
| w = torch.softmax(self.gate_logits[t], 0) | |
| h = (w[None, :, None] * hv).sum(1) | |
| if t in ("tcr", "vcr") and ego_stats is not None: | |
| h = torch.cat([h, ego_stats], -1) | |
| if t in self.frame_avg_tasks: # frame-level logits, mean over observed frames | |
| hs = self.ln_in(hv_seq) # [B,4,T,2048] | |
| hs = (w[None, :, None, None] * hs).sum(1) # [B,T,2048] | |
| x = hs if t in ("tcr", "vcr") else torch.stack([drv_cat(hs[:, i]) for i in range(T)], 1) | |
| out[f"logits_{t}"] = self.heads[t](x).mean(1) | |
| else: | |
| out[f"logits_{t}"] = self.heads[t](drv_cat(h) if t in ("dbr", "der") else h) | |
| return out | |
| z_in = vis[:, :, P.VIEWS.index("incar")] + self.view_emb[P.VIEWS.index("incar")] + self.pos[:T] | |
| ext_idx = [P.VIEWS.index(x) for x in ("front", "left", "right")] | |
| z_ext = (vis[:, :, ext_idx] + self.view_emb[ext_idx]).mean(2) + self.pos[:T] | |
| h_in_seq = self.ln_in(self.enc_in(z_in, src_mask=mask)) # [B,T,2048] causal states | |
| h_in = h_in_seq[:, -1] | |
| h_ext = self.ln_ext(self.enc_ext(z_ext, src_mask=mask)[:, -1]) | |
| if ego_stats is not None: | |
| h_ext = torch.cat([h_ext, ego_stats], -1) | |
| for t in P.TASKS: | |
| if t in ("dbr", "der") and t in self.frame_avg_tasks: | |
| out[f"logits_{t}"] = self.heads[t](torch.stack([drv_cat(h_in_seq[:, i]) for i in range(T)], 1)).mean(1) | |
| else: | |
| out[f"logits_{t}"] = self.heads[t](drv_cat(h_in) if t in ("dbr", "der") else h_ext) | |
| return out | |
| class SemOnly(nn.Module): | |
| """Semantic-only reference: frozen Qwen features -> SemAttnHeads (repro head recipe), no trunk, no geometry | |
| loss. `pred` is the zero-velocity copy so the harness contract holds (its MPJPE == ZV, not a forecast).""" | |
| name = "sem_only" | |
| def __init__(self, use_ego=False, view_gate=False, frame_avg_tasks=()): | |
| super().__init__() | |
| self.use_visual = self.use_semantic = True | |
| self.use_ego = use_ego | |
| self.sem = SemAttnHeads(0, use_ego=use_ego, view_gate=view_gate, frame_avg_tasks=frame_avg_tasks) | |
| if use_ego: | |
| self.register_buffer("ego_mean", torch.zeros(EGO_DIM)) | |
| self.register_buffer("ego_std", torch.ones(EGO_DIM)) | |
| def set_ego_stats(self, mean, std): | |
| self.ego_mean.copy_(torch.as_tensor(mean, dtype=self.ego_mean.dtype)) | |
| self.ego_std.copy_(torch.as_tensor(std, dtype=self.ego_std.dtype).clamp_min(1e-6)) | |
| def ego_norm(self, ego): | |
| return (ego.flatten(2).to(self.ego_mean.dtype) - self.ego_mean) / self.ego_std | |
| def forward(self, obs, vis=None, ego=None): | |
| assert obs.shape[1] == P.T_OBS and vis is not None and vis.shape[1] == P.T_OBS | |
| es = self.ego_norm(ego).mean(1) if (self.use_ego and ego is not None) else None | |
| out = {"pred": obs[:, -1:, :, :2].expand(-1, P.T_PRED, -1, -1), "last": obs[:, -1:, :, :2], "log_b": None} | |
| out.update(self.sem(vis.float(), None, es)) | |
| return out | |
| class CWM3D(nn.Module): | |
| name = "cwm3d" | |
| def __init__(self, d=256, layers=4, heads=4, ffn=1024, head_hidden=512, dropout=0.0, | |
| anchor=True, uncertainty=True, dct=False, rollout="chunked", | |
| use_visual=False, traffic_gate=False, modality_dropout=0.3, | |
| use_semantic=False, sem_input="tokens", use_ego=False, sem_view_gate=False, sem_frame_avg=(), | |
| refine_late=False, formulation="residual", horizon_gain=None): | |
| super().__init__() | |
| assert rollout in ("chunked", "recursive") | |
| assert formulation in ("residual", "abs"), formulation | |
| self.formulation = formulation | |
| if formulation == "abs": | |
| anchor = False # absolute-coordinate read-out: no anchor mixture (decisions 2026-09-15) | |
| self.d, self.rollout, self.uncertainty = d, rollout, uncertainty | |
| self.refine_late = refine_late | |
| self.use_visual, self.traffic_gate, self.use_semantic = use_visual, traffic_gate, use_semantic | |
| assert sem_input in ("tokens", "direct", "attn") | |
| self.sem_input = sem_input | |
| self.modality_dropout = modality_dropout | |
| self.tok = PartTokenizer(d) | |
| self.n_part = len(P.PARTS) | |
| self.use_ego = use_ego | |
| extra = (N_VIEWS if use_visual else 0) + (1 if use_ego else 0) | |
| if use_ego: | |
| # one ego-motion token per observed frame; inputs standardised with train statistics (buffers) | |
| self.ego_in = nn.Linear(EGO_DIM, d) | |
| self.ego_time_emb = nn.Parameter(torch.zeros(P.T_OBS, d)) | |
| nn.init.normal_(self.ego_time_emb, std=0.02) | |
| self.register_buffer("ego_mean", torch.zeros(EGO_DIM)) | |
| self.register_buffer("ego_std", torch.ones(EGO_DIM)) | |
| if use_visual: | |
| self.vis_in = nn.Sequential(nn.LayerNorm(VIS_DIM), nn.Linear(VIS_DIM, d)) | |
| self.view_emb = nn.Parameter(torch.zeros(N_VIEWS, d)) | |
| nn.init.normal_(self.view_emb, std=0.02) | |
| self.vis_time_emb = nn.Parameter(torch.zeros(P.T_OBS, d)) | |
| nn.init.normal_(self.vis_time_emb, std=0.02) | |
| self.blocks = nn.ModuleList(Block(d, heads, ffn, dropout) for _ in range(layers)) | |
| self.ln_f = nn.LayerNorm(d) | |
| self.register_buffer("mask", causal_part_mask(P.T_OBS, self.n_part, extra)) | |
| self.gate = TrafficGate(d, heads) if (use_visual and traffic_gate) else None | |
| self.head = HorizonHead(d, head_hidden, P.T_PRED, dct=dct, zero_init=(formulation == "residual")) | |
| self.mix = AnchorMixture(enabled=anchor) if formulation == "residual" else AbsoluteReadout() | |
| # pre-registered HM h=5 arms (b) / (c): learned per-horizon gain on the anchored residual; None / "none" = unchanged model (no module) | |
| assert horizon_gain in (None, "none") + HorizonGain.MODES, horizon_gain | |
| if horizon_gain in HorizonGain.MODES: | |
| assert formulation == "residual", "horizon_gain needs the residual read-out (there is no residual in the abs formulation)" | |
| self.hgain = HorizonGain(horizon_gain) | |
| else: | |
| self.hgain = None | |
| if refine_late: | |
| # late-horizon refinement (roadmap B3): Δ for h = 4, 5 from [last-frame part tokens, DETACHED h <= 3 | |
| # predictions relative to the last pose / SCALE]; output layer zero-initialised -> identity at init | |
| self.n_early = P.T_PRED - 2 | |
| r_in = self.n_part * d + self.n_early * P.J * 2 | |
| self.refine = nn.Sequential(nn.LayerNorm(r_in), nn.Linear(r_in, head_hidden), nn.GELU(), | |
| nn.Linear(head_hidden, 2 * P.J * 2)) | |
| nn.init.zeros_(self.refine[-1].weight) | |
| nn.init.zeros_(self.refine[-1].bias) | |
| if use_semantic: | |
| if sem_input == "attn" and use_visual: | |
| self.sem = SemAttnHeads(d, use_ego=use_ego, view_gate=sem_view_gate, frame_avg_tasks=sem_frame_avg) | |
| elif sem_input == "direct" and use_visual: | |
| # heads read the frozen 2048-d features directly (mean over the observed window): DBR/DER on | |
| # [LN(in-cabin feature), driver part-token mean], TCR/VCR on LN(mean exterior feature) | |
| self.sem_ln_in, self.sem_ln_ext = nn.LayerNorm(VIS_DIM), nn.LayerNorm(VIS_DIM) | |
| self.sem = nn.ModuleDict({t: nn.Linear(VIS_DIM + (d if t in ("dbr", "der") else (EGO_DIM if use_ego else 0)), | |
| len(P.CLASSES[t])) for t in P.TASKS}) | |
| else: | |
| self.sem = nn.ModuleDict({t: nn.Sequential(nn.LayerNorm(d), nn.Linear(d, len(P.CLASSES[t]))) | |
| for t in P.TASKS}) | |
| # ------------------------------------------------------------------ encoder | |
| def set_ego_stats(self, mean, std): | |
| self.ego_mean.copy_(torch.as_tensor(mean, dtype=self.ego_mean.dtype)) | |
| self.ego_std.copy_(torch.as_tensor(std, dtype=self.ego_std.dtype).clamp_min(1e-6)) | |
| def ego_norm(self, ego): # [B,T,3,16] -> [B,T,48] standardised | |
| return (ego.flatten(2).to(self.ego_mean.dtype) - self.ego_mean) / self.ego_std | |
| def encode(self, obs, vis=None, ego=None, ext_scale=None, gate_override=None): | |
| """-> part tokens [B,T,4,d] (post-transformer), visual tokens [B,T,V,d] or None. | |
| ext_scale: test-time intervention, multiplies the exterior-view visual tokens (front/left/right) by lambda | |
| BEFORE the transformer (None or 1.0 = untouched: no op is executed, so the output is bit-identical); | |
| gate_override: clamps the traffic gate (variant iii) to a constant.""" | |
| B, T = obs.shape[:2] | |
| x = self.tok(obs) # [B,T,4,d] | |
| vtok = None | |
| if self.use_visual: | |
| if vis is None: | |
| vis = obs.new_zeros(B, T, N_VIEWS, VIS_DIM) | |
| vtok = self.vis_in(vis.to(x.dtype)) + self.view_emb[None, None] + self.vis_time_emb[None, :T, None] # feats are fp16 | |
| if self.training and self.modality_dropout > 0: | |
| keep = (torch.rand(B, 1, 1, 1, device=obs.device) > self.modality_dropout).to(vtok.dtype) | |
| vtok = vtok * keep | |
| if ext_scale is not None and float(ext_scale) != 1.0: | |
| sc = torch.ones(N_VIEWS, device=vtok.device, dtype=vtok.dtype) | |
| sc[EXT_VIEWS] = float(ext_scale) | |
| vtok = vtok * sc[None, None, :, None] | |
| x = torch.cat([x, vtok], 2) # [B,T,4+V,d] | |
| if self.use_ego: | |
| if ego is None: | |
| ego = obs.new_zeros(B, T, 3, 16) | |
| etok = self.ego_in(self.ego_norm(ego)) + self.ego_time_emb[None, :T] | |
| x = torch.cat([x, etok[:, :, None]], 2) # [B,T,4+V+1,d] | |
| L = x.shape[1] * x.shape[2] | |
| x = x.reshape(B, L, self.d) | |
| mask = self.mask[:L, :L] | |
| for blk in self.blocks: | |
| x = blk(x, mask) | |
| x = self.ln_f(x).view(B, T, -1, self.d) | |
| parts = x[:, :, :self.n_part] | |
| vtok = x[:, :, self.n_part:self.n_part + N_VIEWS] if self.use_visual else None | |
| if self.gate is not None: | |
| # driver tokens of the last frame attend to external-view tokens of all observed frames | |
| ext = vtok[:, :, EXT_VIEWS].reshape(B, -1, self.d) | |
| q = self.gate(parts[:, -1], ext, None, override=gate_override) | |
| parts = torch.cat([parts[:, :-1], q[:, None]], 1) | |
| return parts, vtok | |
| def _decode(self, parts, obs): | |
| anc = anchors(obs) | |
| ho = self.head(parts[:, -1]) # [B,H,J,7] | |
| if not self.uncertainty: | |
| ho = torch.cat([ho[..., :2], torch.zeros_like(ho[..., 2:3]), ho[..., 3:]], -1) | |
| pred, log_b = self.mix(ho, anc) | |
| if self.hgain is not None: # arms (b)/(c): pred = last + g_h ⊙ (pred − last) | |
| pred = self.hgain(pred, anc[:, 0:1]) | |
| return pred, log_b, ho | |
| # ------------------------------------------------------------------ forward | |
| def forward(self, obs, vis=None, ego=None, ext_scale=None, gate_override=None): | |
| assert obs.shape[1] == P.T_OBS and obs.shape[2] == P.J, obs.shape | |
| assert ego is None or ego.shape[1] == P.T_OBS, ego.shape | |
| out = {"last": obs[:, -1:, :, :2]} | |
| if self.rollout == "chunked": | |
| parts, vtok = self.encode(obs, vis, ego, ext_scale=ext_scale, gate_override=gate_override) | |
| pred, log_b, ho = self._decode(parts, obs) | |
| else: | |
| win, preds, logbs = obs, [], [] | |
| parts = vtok = None | |
| for _ in range(P.T_PRED): | |
| parts, vtok = self.encode(win, vis, ego, ext_scale=ext_scale, gate_override=gate_override) | |
| p, lb, ho = self._decode(parts, win) | |
| preds.append(p[:, 0]); logbs.append(lb[:, 0]) | |
| nxt = torch.cat([p[:, 0], win[:, -1, :, 2:3]], -1) # carry last observed conf | |
| win = torch.cat([win[:, 1:], nxt[:, None]], 1) | |
| pred, log_b = torch.stack(preds, 1), torch.stack(logbs, 1) | |
| if self.refine_late: | |
| last = obs[:, -1:, :, :2] | |
| early = ((pred[:, :self.n_early] - last) / SCALE).detach().flatten(1) # [B, 3*J*2], no gradient | |
| r_in = torch.cat([parts[:, -1].flatten(1), early], 1) | |
| delta = self.refine(r_in).view(-1, 2, P.J, 2) | |
| pred = torch.cat([pred[:, :self.n_early], pred[:, self.n_early:] + SCALE * delta], 1) | |
| out["refine_delta_px"] = (SCALE * delta).detach() | |
| out["pred"] = pred | |
| out["log_b"] = log_b if self.uncertainty else None | |
| if self.mix.enabled: | |
| out["anchor_w"] = torch.softmax(ho[..., HorizonHead.C_LOGIT], -1) | |
| out["anchor_gate"] = ho[..., HorizonHead.C_GATE] | |
| if self.use_semantic and self.sem_input == "attn" and self.use_visual: | |
| v = (obs.new_zeros(obs.shape[0], obs.shape[1], N_VIEWS, VIS_DIM) if vis is None else vis).to(parts.dtype) | |
| es = None | |
| if self.use_ego: | |
| es = self.ego_norm(obs.new_zeros(obs.shape[0], obs.shape[1], 3, 16) if ego is None else ego).mean(1) | |
| out.update(self.sem(v, parts.mean((1, 2)), es)) | |
| elif self.use_semantic and self.sem_input == "direct" and self.use_visual: | |
| v = (obs.new_zeros(obs.shape[0], obs.shape[1], N_VIEWS, VIS_DIM) if vis is None else vis).to(parts.dtype) | |
| f_in = self.sem_ln_in(v[:, :, P.VIEWS.index("incar")].mean(1)) | |
| f_ext = self.sem_ln_ext(v[:, :, [P.VIEWS.index(x) for x in ("front", "left", "right")]].mean((1, 2))) | |
| drv = parts.mean((1, 2)) | |
| if self.use_ego: | |
| e = self.ego_norm(obs.new_zeros(obs.shape[0], obs.shape[1], 3, 16) if ego is None else ego).mean(1) | |
| f_ext = torch.cat([f_ext, e], -1) | |
| for t in P.TASKS: | |
| out[f"logits_{t}"] = self.sem[t](torch.cat([f_in, drv], -1) if t in ("dbr", "der") else f_ext) | |
| elif self.use_semantic: | |
| drv = parts.mean((1, 2)) # driver parts, observed window only | |
| if vtok is not None: | |
| incar = vtok[:, :, P.VIEWS.index("incar")].mean(1) | |
| ext = vtok[:, :, [P.VIEWS.index(v) for v in ("front", "left", "right")]].mean((1, 2)) | |
| drv_feat, ext_feat = drv + incar, ext | |
| else: | |
| drv_feat, ext_feat = drv, drv | |
| for t in P.TASKS: | |
| out[f"logits_{t}"] = self.sem[t](drv_feat if t in ("dbr", "der") else ext_feat) | |
| return out | |
| def build(name: str, **kw) -> nn.Module: | |
| """cwm3d, cwm3d_noanchor, cwm3d_nounc, cwm3d_recursive, cwm3d_dct (physics is a loss flag).""" | |
| base = dict(kw) | |
| if name == "cwm3d": | |
| return CWM3D(**base) | |
| if name == "cwm3d_noanchor": | |
| return CWM3D(**{**base, "anchor": False}) | |
| if name == "cwm3d_nounc": | |
| return CWM3D(**{**base, "uncertainty": False}) | |
| if name == "cwm3d_recursive": | |
| return CWM3D(**{**base, "rollout": "recursive"}) | |
| if name == "cwm3d_dct": | |
| return CWM3D(**{**base, "dct": True}) | |
| if name == "cwm3d_vis": | |
| return CWM3D(**{**base, "use_visual": True, "traffic_gate": True}) | |
| if name == "sem_only": | |
| return SemOnly(use_ego=base.get("use_ego", False), view_gate=base.get("sem_view_gate", False), | |
| frame_avg_tasks=base.get("sem_frame_avg", ())) | |
| raise ValueError(name) | |
| if __name__ == "__main__": | |
| m = CWM3D() | |
| print("params", 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]) | |
| o = m(x) | |
| print({k: tuple(v.shape) for k, v in o.items() if torch.is_tensor(v)}) | |
| print("zero-init == ZV:", torch.allclose(o["pred"], x[:, -1:, :, :2].expand_as(o["pred"]))) | |
Xet Storage Details
- Size:
- 22.1 kB
- Xet hash:
- e6b179c04e2f1052da21d41eb05bc10b05428a1469dfa133bd38e6ca409c5261
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.