| |
| """Pure-PyTorch fp32 re-implementation of the three deployed StreamPETR graphs with per-module taps. |
| |
| Ported from the verified research reference (``research/streampetr/scripts/torch_reference.py``: PCC >= 0.9999999998 |
| against ONNX Runtime on every output, SPEC section 6) onto this port's weight loader (``reference.weights``: every |
| parameter and constant addressed by its consuming ONNX node). Module boundaries follow the training code (AWML |
| ``projects/StreamPETR``: VoVNet-99-eSE + CPFPN, PETR position embedding, StreamPETRHead with the PETRTemporal |
| decoder) so that each TT stage has a reference tap. No mmdet / mmcv / mmdet3d; plain ``torch.nn.functional``. |
| |
| ``Backbone(img)``: ``img`` [1, 5, 3, 480, 640] (normalised, ``host.preprocess``) -> ``img_feats`` [1, 5, 256, 30, 40]. |
| ``PositionEmbedding(img_metas_pad, intrinsics, img2lidar)`` -> ``pos_embed`` [1, 6000, 256], ``cone`` [1, 6000, 8]. |
| ``Head(x, pos_embed, cone, E, Einv, pre)``: the ONNX head (global-frame memory) -> its 14 outputs (ONNX names) plus |
| ``topk_idx`` [1, 256, 1] and ``rec_score`` [1, 900, 1] (not ONNX outputs). The stages are separate methods |
| (``tokens``, ``memory_pre_update``, ``temporal``, ``queries``, ``decoder``, ``predict``, ``post_update``) so that |
| ``reference.port_form`` can run the TT port's re-parameterised memory through the same functions. |
| |
| Tap names (``ttaw.golden.TapRegistry``; the device PCC gates use the same names): |
| |
| ================================ ====================== =================================================== |
| tap shape where |
| ================================ ====================== =================================================== |
| ``backbone.stem`` [5, 128, 120, 160] stem_1 (s2) .. stem_3 (s2), conv + ReLU |
| ``backbone.stage2.osa1`` ... [5, C, h, w] every OSA block output (after eSE and the identity) |
| ``backbone.stage2`` .. ``stage5`` [5, 256..1024, ...] stage outputs (stage4 = 30x40, stage5 = 15x20) |
| ``neck.lateral0`` / ``lateral1`` [5, 256, 30, 40] / 15x20 CPFPN 1x1 laterals (no activation) |
| ``neck.merged`` [5, 256, 30, 40] lateral0 + nearest x2(lateral1) |
| ``img_feats`` [1, 5, 256, 30, 40] fpn_convs.0 3x3 (the backbone graph output) |
| ``pe.coords`` [1, 6000, 192] normalised 3-D points of the 64 depth bins |
| ``pe.hidden`` [1, 6000, 1024] ReLU(Linear 192 -> 1024) of inverse_sigmoid(coords) |
| ``pos_embed`` / ``cone`` [1, 6000, 256] / 8 the PE graph outputs (once per calibration) |
| ``head.tokens`` [1, 6000, 256] x permuted to camera-major tokens |
| ``head.memory_embed`` [1, 6000, 256] memory_embed MLP |
| ``head.sa_gamma`` / ``sa_beta`` [1, 6000, 256] spatial-alignment MLN coefficients (cone only) |
| ``head.memory`` [1, 6000, 256] LN(memory) * gamma + beta: the cross-attention values |
| ``head.pos`` [1, 6000, 256] pos_embed * sigmoid(featurized PE): + memory = keys |
| ``head.mem_ego`` / ``mem_ref`` [1, 1024, 4, 4] / 3 Einv @ pre_memory_egopose / reference points |
| ``head.temp_ref`` [1, 1024, 3] (mem_ref - pc_min) / pc_size |
| ``head.motion`` [1, 1024, 180] nerf_pe([velo, ts, mem_ego[:3]]) |
| ``head.temp_pos`` / ``temp_mem`` [1, 1024, 256] after the two MLNs and the time embedding |
| ``head.tgt`` / ``query_pos`` [1, 900, 256] decoder inputs (644 learned + 256 propagated) |
| ``head.reference_points`` [1, 900, 3] |
| ``head.layer<l>.self_attn`` [1, 900, 256] LN(tgt + self-attention), l = 0..5 |
| ``head.layer<l>.cross_attn`` [1, 900, 256] LN(. + cross-attention over the 6000 tokens) |
| ``head.layer<l>.ffn`` [1, 900, 256] LN(. + FFN): the next layer's tgt |
| ``head.outs_dec`` [6, 1, 900, 256] post_norm of every layer's tgt |
| ``head.cls`` / ``head.reg`` [6, 1, 900, 5] / 10 shared branches; reg centres decoded to metres |
| ``head.rec_score`` [1, 900, 1] max_c sigmoid(cls[5]) (TopK k=1) |
| ``head.topk_idx`` [1, 256, 1] TopK k=256 (sorted, ties to the lower index) |
| ``all_cls_scores`` ... ONNX outputs ``Head.__call__`` return values (ONNX names) |
| ================================ ====================== =================================================== |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any, Dict, Optional, Tuple |
|
|
| import numpy as np |
|
|
| from ..ttaw.golden import NULL_TAPS, TapRegistry |
| from .config import StreamPETRConfig, default_config |
| from .weights import OSA_STAGES, StreamPETRWeights |
|
|
| __all__ = ["Backbone", "PositionEmbedding", "Head", "StreamPETRNet", "inverse_sigmoid", "nerf_pe", "posemb", |
| "HEAD_OUTPUTS", "MEMORY_KEYS"] |
|
|
| HEAD_OUTPUTS = ("all_cls_scores", "all_bbox_preds", "post_memory_embedding", "post_memory_reference_point", |
| "post_memory_timestamp", "post_memory_egopose", "post_memory_velo", "reference_points", "tgt", |
| "temp_memory", "temp_pos", "query_pos", "query_pos_in", "outs_dec") |
| MEMORY_KEYS = ("embedding", "reference_point", "timestamp", "egopose", "velo") |
|
|
|
|
| def _torch(): |
| import torch |
|
|
| return torch |
|
|
|
|
| def _t(a: Any): |
| """numpy / torch -> float32 torch tensor (own copy for numpy input).""" |
| torch = _torch() |
| if isinstance(a, torch.Tensor): |
| return a.float() |
| return torch.from_numpy(np.array(a, dtype=np.float32)) |
|
|
|
|
| def inverse_sigmoid(x: Any, eps: float = 1e-4): |
| """``misc.py:7-25`` as the ONNX graphs compute it: clip(x, 0, 1); log(max(x, eps) / max(1 - x, eps)).""" |
| x = x.clamp(0, 1) |
| return (x.clamp(min=eps) / (1 - x).clamp(min=eps)).log() |
|
|
|
|
| def nerf_pe(x: Any, n: int = 6): |
| """[sin(x), cos(x), sin(2x), cos(2x), ..., sin(32x), cos(32x)] along the last axis (ONNX /Concat_20).""" |
| torch = _torch() |
| enc = [torch.sin(x), torch.cos(x)] |
| for f in (2.0, 4.0, 8.0, 16.0, 32.0)[: n - 1]: |
| xf = x * f |
| enc += [torch.sin(xf), torch.cos(xf)] |
| return torch.cat(enc, -1) |
|
|
|
|
| def posemb(p: Any, dim_t: Any): |
| """One axis of ``pos2posemb``: ``q = p / dim_t``, then (sin of the even, cos of the odd entries) interleaved.""" |
| torch = _torch() |
| q = p[..., None] / dim_t |
| return torch.stack((q[..., 0::2].sin(), q[..., 1::2].cos()), -1).flatten(-2) |
|
|
|
|
| |
|
|
| class Backbone: |
| """VoVNet-99-eSE + CPFPN (``vovnet.py:80-88,174-239``, ``cp_fpn.py:160-204``), BN folded by the exporter.""" |
|
|
| def __init__(self, weights: StreamPETRWeights, cfg: Optional[StreamPETRConfig] = None): |
| self.cfg = cfg or default_config() |
| self.params: Dict[str, Tuple[Any, Any, int, int]] = { |
| n: (_t(c.weight), None if c.bias is None else _t(c.bias), c.stride, c.padding) |
| for n, c in weights.iter_convs()} |
|
|
| def conv(self, x: Any, node: str, relu: bool = True): |
| import torch.nn.functional as F |
|
|
| w, b, s, p = self.params[node] |
| y = F.conv2d(x, w, b, stride=s, padding=p) |
| return F.relu(y) if relu else y |
|
|
| def osa(self, x: Any, stage: int, k: int, identity: bool): |
| """One OSA block: 5 chained 3x3 convs, concat [x, 5 outputs], 1x1 concat conv, eSE, optional identity.""" |
| torch = _torch() |
| pre = f"/img_backbone/stage{stage}/OSA{stage}_{k}" |
| outs, h = [x], x |
| for i in range(5): |
| h = self.conv(h, f"{pre}/layers.{i}/OSA{stage}_{k}_{i}/conv/Conv") |
| outs.append(h) |
| xt = self.conv(torch.cat(outs, 1), f"{pre}/concat/OSA{stage}_{k}_concat/conv/Conv") |
| |
| s = self.conv(xt.mean((2, 3), keepdim=True), f"{pre}/ese/fc/Conv", relu=False) |
| xt = xt * (torch.clamp(s + 3.0, 0.0, 6.0) / 6.0) |
| return xt + x if identity else xt |
|
|
| def __call__(self, img: Any, taps: TapRegistry = NULL_TAPS): |
| """``img`` [1, N, 3, H, W] (or [N, 3, H, W]) normalised -> ``img_feats`` [1, N, 256, H/16, W/16].""" |
| import torch.nn.functional as F |
|
|
| x = _t(img) |
| if x.dim() == 4: |
| x = x[None] |
| b, n = x.shape[:2] |
| x = x.reshape(b * n, *x.shape[2:]) |
| for i in (1, 2, 3): |
| x = self.conv(x, f"/img_backbone/stem/stem_{i}/conv/Conv") |
| taps.tap("backbone.stem", x) |
| feats = {} |
| for st, (nb, _, _) in OSA_STAGES.items(): |
| if st != 2: |
| x = F.max_pool2d(x, 3, 2, ceil_mode=True) |
| for k in range(1, nb + 1): |
| x = taps.tap(f"backbone.stage{st}.osa{k}", self.osa(x, st, k, identity=k > 1)) |
| feats[st] = taps.tap(f"backbone.stage{st}", x) |
| lat0 = taps.tap("neck.lateral0", self.conv(feats[4], "/img_neck/lateral_convs.0/conv/Conv", relu=False)) |
| lat1 = taps.tap("neck.lateral1", self.conv(feats[5], "/img_neck/lateral_convs.1/conv/Conv", relu=False)) |
| |
| merged = taps.tap("neck.merged", lat0 + F.interpolate(lat1, size=lat0.shape[-2:], mode="nearest")) |
| out = self.conv(merged, "/img_neck/fpn_convs.0/conv/Conv", relu=False) |
| return taps.tap("img_feats", out.reshape(b, n, *out.shape[1:])) |
|
|
|
|
| |
|
|
| class PositionEmbedding: |
| """PETR 3-D position embedding (``streampetr_head.py:415-463``, ``containers.py:11-24``), run once per |
| calibration. Constants (pixel centres, LID depth bins, ``position_range``) are read from the ONNX graph.""" |
|
|
| def __init__(self, weights: StreamPETRWeights, cfg: Optional[StreamPETRConfig] = None): |
| self.cfg = cfg or default_config() |
| pe = weights.pe |
| self.l0 = tuple(_t(a) for a in weights.linear("/position_encoder/position_encoder.0", "position_embedding")) |
| self.l2 = tuple(_t(a) for a in weights.linear("/position_encoder/position_encoder.2", "position_embedding")) |
| self.u = _t(pe.params("/Div")[0]) |
| self.v = _t(pe.params("/Div_1")[0]) |
| depth = np.asarray(pe.array("/Tile_3_output_0"), np.float32) |
| if not np.all(depth == depth[:, :1]): |
| raise ValueError("the PE depth-bin constant differs between tokens") |
| self.depth = _t(depth[0, 0, :, 0]) |
| self.depth_min = float(np.asarray(pe.array("/Constant_71_output_0")).reshape(-1)[0]) |
| self.pr_min = _t(pe.params("/Sub")[1]) |
| self.pr_size = _t(pe.params("/Div_3")[1]) |
| self.focal_div = float(pe.params("/Div_2")[1]) |
|
|
| def __call__(self, img_metas_pad: Any, intrinsics: Any, img2lidar: Any, taps: TapRegistry = NULL_TAPS): |
| """``img_metas_pad`` [3] = (pad_h, pad_w, 3); ``intrinsics`` / ``img2lidar`` [1, N, 4, 4] float32 -> |
| ``pos_embed`` [1, N*H*W, 256], ``cone`` [1, N*H*W, 8].""" |
| torch = _torch() |
| metas, k4, i2l = _t(img_metas_pad), _t(intrinsics), _t(img2lidar) |
| h, w = self.cfg.feat_hw |
| n = k4.shape[1] |
| pad_h, pad_w = metas[0], metas[1] |
| xs = (self.u / pad_w) * pad_w |
| ys = (self.v / pad_h) * pad_h |
| sy, sx = torch.meshgrid(ys, xs, indexing="ij") |
| loc = torch.stack([sx, sy], -1) |
| centers = loc[None].repeat(n, 1, 1, 1).reshape(1, n * h * w, 1, 2) |
| dn = self.depth.numel() |
| d = self.depth.view(1, 1, dn, 1).expand(1, n * h * w, dn, 1) |
| coords = torch.cat([centers.expand(1, n * h * w, dn, 2), d, torch.ones_like(d)], -1) |
| coords = torch.cat([coords[..., :2] * torch.clamp(coords[..., 2:3], min=self.depth_min), coords[..., 2:]], -1) |
| m = i2l.reshape(n, 1, 1, 4, 4).expand(n, h * w, dn, 4, 4).reshape(1, n * h * w, dn, 4, 4) |
| c3d = (m @ coords.unsqueeze(-1)).squeeze(-1)[..., :3] |
| c3d = ((c3d - self.pr_min) / self.pr_size).reshape(1, n * h * w, dn * 3) |
| taps.tap("pe.coords", c3d) |
| hid = taps.tap("pe.hidden", torch.relu(inverse_sigmoid(c3d) @ self.l0[0] + self.l0[1])) |
| pos = hid @ self.l2[0] + self.l2[1] |
| |
| f = torch.stack([k4[..., 0, 0], k4[..., 1, 1]], -1).abs() / self.focal_div |
| f = f.repeat(1, h * w, 1) |
| cone = torch.cat([f, c3d[..., -3:], c3d[..., -90:-87]], -1) |
| return taps.tap("pos_embed", pos), taps.tap("cone", cone) |
|
|
|
|
| |
|
|
| class Head: |
| """StreamPETRHead + PETRTemporalTransformer decoder exactly as ``simplify_pts_head_memory.onnx`` (S:4.3).""" |
|
|
| def __init__(self, weights: StreamPETRWeights, cfg: Optional[StreamPETRConfig] = None): |
| self.cfg = cfg or default_config() |
| lin = lambda p: tuple(_t(a) for a in weights.linear(p)) |
| ln = lambda n: (lambda g, b, e: (_t(g), _t(b), e))(*weights.layer_norm(n)) |
| self.memory_embed = (lin("/memory_embed/memory_embed.0"), lin("/memory_embed/memory_embed.2")) |
| self.sa = dict(ln=ln("/spatial_alignment/ln/LayerNormalization"), reduce=lin("/spatial_alignment/reduce/reduce.0"), |
| gamma=lin("/spatial_alignment/gamma"), beta=lin("/spatial_alignment/beta")) |
| self.fpe = (lin("/featurized_pe/conv_reduce"), lin("/featurized_pe/conv_expand")) |
| self.qemb = (lin("/query_embedding/query_embedding.0_1"), lin("/query_embedding/query_embedding.2_1")) |
| self.mln_pe = dict(ln=ln("/ego_pose_pe/ln_1/LayerNormalization"), reduce=lin("/ego_pose_pe/reduce/reduce.0_1"), |
| gamma=lin("/ego_pose_pe/gamma_1"), beta=lin("/ego_pose_pe/beta_1")) |
| self.mln_mem = dict(ln=ln("/ego_pose_memory/ln_1/LayerNormalization"), |
| reduce=lin("/ego_pose_memory/reduce/reduce.0_1"), |
| gamma=lin("/ego_pose_memory/gamma_1"), beta=lin("/ego_pose_memory/beta_1")) |
| self.time = (lin("/time_embedding/time_embedding.0_1"), |
| ln("/time_embedding/time_embedding.1_1/LayerNormalization")) |
| self.scale = float(weights.attention_scale()) |
| self.layers = [] |
| for li in range(self.cfg.num_layers): |
| p = f"/transformer/decoder/layers.{li}" |
| attn = [] |
| for a in range(2): |
| q = f"{p}/attentions.{a}" |
| attn.append({k: tuple(_t(v) for v in weights.attn_proj(q, k)) for k in ("q", "k", "v")} |
| | {"o": tuple(_t(v) for v in weights.gemm(f"{q}/attn/Gemm"))}) |
| self.layers.append(dict(attn=attn, norms=[ln(f"{p}/norms.{i}/LayerNormalization") for i in range(3)], |
| ffn=(lin(f"{p}/ffns.0/layers/layers.0/layers.0.0"), |
| lin(f"{p}/ffns.0/layers/layers.1")))) |
| self.post_norm = ln("/transformer/decoder/post_norm/LayerNormalization") |
| self.cls = (lin("/cls_branches.0/cls_branches.0.0"), ln("/cls_branches.0/cls_branches.0.1/LayerNormalization"), |
| lin("/cls_branches.0/cls_branches.0.3"), ln("/cls_branches.0/cls_branches.0.4/LayerNormalization"), |
| lin("/cls_branches.0/cls_branches.0.6")) |
| self.reg = (lin("/reg_branches.0/reg_branches.0.0"), lin("/reg_branches.0/reg_branches.0.2"), |
| lin("/reg_branches.0/reg_branches.0.4")) |
| q = weights.learned_queries() |
| self.ref644, self.tgt644 = _t(q["reference_points"]), _t(q["tgt"]) |
| self.qpos644, self.query_pos_in = _t(q["query_pos"]), _t(q["query_pos_in"]) |
| h = weights.head |
| self.pc_min = _t(h.params("/Sub")[1]) |
| self.pc_size = _t(h.params("/Div_3")[1]) |
| self.two_pi = float(h.params("/Mul_1")[1]) |
| self.dim_t3 = _t(h.params("/Div_4")[1]) |
| self.dim_t1 = _t(h.params("/Div_7")[1]) |
| self.box_scale = _t(h.params("/Mul_69")[1]) |
| self.box_min = _t(h.params("/Add_8")[1]) |
|
|
| |
| @staticmethod |
| def lin(x: Any, wb: Tuple[Any, Any]): |
| return x @ wb[0] + wb[1] |
|
|
| @staticmethod |
| def layer_norm(x: Any, p: Tuple[Any, Any, float]): |
| import torch.nn.functional as F |
|
|
| return F.layer_norm(x, (x.shape[-1],), p[0], p[1], p[2]) |
|
|
| def mln(self, x: Any, c: Any, p: Dict[str, Any]): |
| """MLN (``petr_transformer.py`` / ``misc.py:155-216``): gamma(c) * LN(x) + beta(c), c = ReLU(reduce(c)).""" |
| torch = _torch() |
| x = self.layer_norm(x, p["ln"]) |
| c = torch.relu(self.lin(c, p["reduce"])) |
| return self.lin(c, p["gamma"]) * x + self.lin(c, p["beta"]) |
|
|
| def mha(self, q_in: Any, k_in: Any, v_in: Any, p: Dict[str, Any]): |
| """nn.MultiheadAttention as exported: Q scaled after the projection, softmax over keys, Gemm out-proj.""" |
| torch = _torch() |
| lq, lk, c = q_in.shape[1], k_in.shape[1], q_in.shape[2] |
| nh = self.cfg.num_heads |
| hd = c // nh |
| q = self.lin(q_in, p["q"]).view(lq, nh, hd).transpose(0, 1) * self.scale |
| k = self.lin(k_in, p["k"]).view(lk, nh, hd).transpose(0, 1) |
| v = self.lin(v_in, p["v"]).view(lk, nh, hd).transpose(0, 1) |
| a = torch.softmax(q @ k.transpose(1, 2), -1) |
| o = (a @ v).transpose(0, 1).reshape(1, lq, c) |
| return self.lin(o, p["o"]) |
|
|
| def posemb3d(self, p: Any): |
| """``pos2posemb3d`` (``positional_encoding.py``): 128 features per axis, concatenated in (y, x, z) order.""" |
| torch = _torch() |
| q = p * self.two_pi |
| x, y, z = (posemb(q[..., i], self.dim_t3) for i in range(3)) |
| return torch.cat((y, x, z), -1) |
|
|
| def posemb1d(self, p: Any): |
| return posemb(p[..., 0] * self.two_pi, self.dim_t1) |
|
|
| def cls_branch(self, x: Any): |
| torch = _torch() |
| x = torch.relu(self.layer_norm(self.lin(x, self.cls[0]), self.cls[1])) |
| x = torch.relu(self.layer_norm(self.lin(x, self.cls[2]), self.cls[3])) |
| return self.lin(x, self.cls[4]) |
|
|
| def reg_branch(self, x: Any): |
| torch = _torch() |
| x = torch.relu(self.lin(x, self.reg[0])) |
| x = torch.relu(self.lin(x, self.reg[1])) |
| return self.lin(x, self.reg[2]) |
|
|
| |
| def spatial_alignment_coeffs(self, cone: Any) -> Tuple[Any, Any]: |
| """gamma / beta of the spatial-alignment MLN: functions of ``cone`` only (constant per calibration).""" |
| torch = _torch() |
| c = torch.relu(self.lin(_t(cone), self.sa["reduce"])) |
| return self.lin(c, self.sa["gamma"]), self.lin(c, self.sa["beta"]) |
|
|
| def tokens(self, x: Any, pos_embed: Any, cone: Any, taps: TapRegistry = NULL_TAPS, |
| sa_coeffs: Optional[Tuple[Any, Any]] = None) -> Tuple[Any, Any]: |
| """Image tokens (``containers.py:76-87``): memory = MLN_sa(memory_embed(x tokens), cone); pos = pos_embed * |
| sigmoid(featurized PE(memory)). ``sa_coeffs`` = precomputed (gamma, beta) (the port's calibration cache).""" |
| torch = _torch() |
| x = _t(x) |
| b, n, c, h, w = x.shape |
| memory = taps.tap("head.tokens", x.permute(0, 1, 3, 4, 2).reshape(b, n * h * w, c)) |
| memory = torch.relu(self.lin(memory, self.memory_embed[0])) |
| memory = taps.tap("head.memory_embed", self.lin(memory, self.memory_embed[1])) |
| gamma, beta = sa_coeffs if sa_coeffs is not None else self.spatial_alignment_coeffs(cone) |
| taps.tap("head.sa_gamma", gamma) |
| taps.tap("head.sa_beta", beta) |
| memory = taps.tap("head.memory", gamma * self.layer_norm(memory, self.sa["ln"]) + beta) |
| gate = torch.sigmoid(self.lin(torch.relu(self.lin(memory, self.fpe[0])), self.fpe[1])) |
| return memory, taps.tap("head.pos", _t(pos_embed) * gate) |
|
|
| @staticmethod |
| def memory_pre_update(pre: Dict[str, Any], Einv: Any, taps: TapRegistry = NULL_TAPS) -> Tuple[Any, Any]: |
| """``containers.py:154-162``: mem_ego = Einv @ pre_egopose; mem_ref = (Einv @ [pre_ref, 1])[:3].""" |
| torch = _torch() |
| einv = _t(Einv) |
| mem_ego = einv[:, None] @ _t(pre["pre_memory_egopose"]) |
| ref = _t(pre["pre_memory_reference_point"]) |
| ref_h = torch.cat([ref, torch.ones_like(ref[..., :1])], -1) |
| mem_ref = (einv[:, None] @ ref_h.unsqueeze(-1)).squeeze(-1)[..., :3] |
| return taps.tap("head.mem_ego", mem_ego), taps.tap("head.mem_ref", mem_ref) |
|
|
| def temporal(self, mem_ref: Any, mem_ego: Any, mem_ts: Any, mem_emb: Any, mem_velo: Any, |
| taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]: |
| """Temporal alignment (``streampetr_head.py:465-507``) of the 1024 memory rows.""" |
| torch = _torch() |
| mem_ts, mem_emb, mem_velo = _t(mem_ts), _t(mem_emb), _t(mem_velo) |
| temp_ref = taps.tap("head.temp_ref", (mem_ref - self.pc_min) / self.pc_size) |
| temp_pos = self.lin(torch.relu(self.lin(self.posemb3d(temp_ref), self.qemb[0])), self.qemb[1]) |
| motion = taps.tap("head.motion", nerf_pe(torch.cat([mem_velo, mem_ts, mem_ego[..., :3, :].flatten(-2)], -1))) |
| temp_pos = self.mln(temp_pos, motion, self.mln_pe) |
| temp_mem = taps.tap("head.temp_mem", self.mln(mem_emb, motion, self.mln_mem)) |
| temp_pos = temp_pos + self.layer_norm(self.lin(self.posemb1d(mem_ts), self.time[0]), self.time[1]) |
| return {"temp_ref": temp_ref, "temp_pos": taps.tap("head.temp_pos", temp_pos), "temp_mem": temp_mem} |
|
|
| def queries(self, t: Dict[str, Any], taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]: |
| """644 learned + the first 256 memory rows; rows 256..1023 feed only the self-attention keys / values.""" |
| torch = _torch() |
| npr = self.cfg.num_propagated |
| tgt = taps.tap("head.tgt", torch.cat([self.tgt644, t["temp_mem"][:, :npr]], 1)) |
| query_pos = taps.tap("head.query_pos", torch.cat([self.qpos644, t["temp_pos"][:, :npr]], 1)) |
| ref = taps.tap("head.reference_points", torch.cat([self.ref644, t["temp_ref"][:, :npr]], 1)) |
| return {"tgt": tgt, "query_pos": query_pos, "reference_points": ref, |
| "temp_memory": t["temp_mem"][:, npr:], "temp_pos": t["temp_pos"][:, npr:]} |
|
|
| def decoder(self, tgt: Any, query_pos: Any, temp_mem: Any, temp_pos: Any, memory: Any, pos: Any, |
| taps: TapRegistry = NULL_TAPS): |
| """6 post-norm layers (``petr_transformer.py:646-753``): self-attention over [queries; memory rows 256..], |
| cross-attention over the 6000 image tokens (keys memory + pos), FFN; returns outs_dec [6, 1, 900, 256].""" |
| torch = _torch() |
| outs = [] |
| for li, lp in enumerate(self.layers): |
| kv = torch.cat([tgt, temp_mem], 1) |
| kpos = torch.cat([query_pos, temp_pos], 1) |
| tgt = self.layer_norm(tgt + self.mha(tgt + query_pos, kv + kpos, kv, lp["attn"][0]), lp["norms"][0]) |
| taps.tap(f"head.layer{li}.self_attn", tgt) |
| tgt = self.layer_norm(tgt + self.mha(tgt + query_pos, memory + pos, memory, lp["attn"][1]), lp["norms"][1]) |
| taps.tap(f"head.layer{li}.cross_attn", tgt) |
| hidden = self.lin(torch.relu(self.lin(tgt, lp["ffn"][0])), lp["ffn"][1]) |
| tgt = taps.tap(f"head.layer{li}.ffn", self.layer_norm(tgt + hidden, lp["norms"][2])) |
| outs.append(self.layer_norm(tgt, self.post_norm)) |
| return taps.tap("head.outs_dec", torch.stack(outs)) |
|
|
| def predict(self, outs_dec: Any, reference_points: Any, taps: TapRegistry = NULL_TAPS) -> Tuple[Any, Any]: |
| """Shared cls / reg branches on every layer; centres = sigmoid(reg[:3] + inverse_sigmoid(ref)) scaled to |
| pc_range (``streampetr_head.py:626-714``). -> cls [6, 1, 900, 5] logits, reg [6, 1, 900, 10].""" |
| torch = _torch() |
| reference = inverse_sigmoid(reference_points) |
| cls = torch.stack([self.cls_branch(o) for o in outs_dec]) |
| reg = [] |
| for o in outs_dec: |
| r = self.reg_branch(o) |
| reg.append(torch.cat([torch.sigmoid(r[..., :3] + reference), r[..., 3:]], -1)) |
| reg = torch.stack(reg) |
| reg = torch.cat([reg[..., :3] * self.box_scale + self.box_min, reg[..., 3:]], -1) |
| return taps.tap("head.cls", cls), taps.tap("head.reg", reg) |
|
|
| def select(self, cls: Any, taps: TapRegistry = NULL_TAPS) -> Tuple[Any, Any]: |
| """``containers.py:42-71``: rec_score = max_c sigmoid(cls[last]) (TopK k=1), then TopK k=256 over the 900 |
| queries (sorted descending, ties to the lower index as ONNX TopK).""" |
| torch = _torch() |
| rec_score = taps.tap("head.rec_score", torch.sigmoid(cls[-1]).max(-1, keepdim=True).values) |
| |
| idx = torch.sort(rec_score[..., 0], dim=1, descending=True, stable=True).indices[:, : self.cfg.num_propagated] |
| return rec_score, taps.tap("head.topk_idx", idx[..., None]) |
|
|
| def post_update(self, cls: Any, reg: Any, outs_dec: Any, mem: Dict[str, Any], E: Any, |
| taps: TapRegistry = NULL_TAPS, topk_idx: Any = None) -> Dict[str, Any]: |
| """``containers.py:42-71``: the new 256 rows (outs_dec / reg of the last layer at the top-256), prepended to |
| the 1024 memory rows; geometry back to the global frame with E. ``topk_idx`` teacher-forces the selection.""" |
| torch = _torch() |
| rec_score, idx = self.select(cls, taps) |
| if topk_idx is not None: |
| idx = torch.as_tensor(np.asarray(topk_idx)).long().reshape(1, -1, 1) |
| g = lambda t: torch.gather(t, 1, idx.expand(-1, -1, t.shape[-1])) |
| rec_ref, rec_mem, rec_velo = g(reg[-1][..., :3]), g(outs_dec[-1]), g(reg[-1][..., -2:]) |
| k = idx.shape[1] |
| rec_ego = torch.eye(4).expand(1, k, 4, 4) |
| e = _t(E) |
| post_ego = e[:, None] @ torch.cat([rec_ego, mem["mem_ego"]], 1) |
| pr = torch.cat([rec_ref, mem["mem_ref"]], 1) |
| post_ref = (e[:, None] @ torch.cat([pr, torch.ones_like(pr[..., :1])], -1).unsqueeze(-1)).squeeze(-1)[..., :3] |
| return {"post_memory_embedding": torch.cat([rec_mem, _t(mem["mem_emb"])], 1), |
| "post_memory_reference_point": post_ref, |
| "post_memory_timestamp": torch.cat([torch.zeros(1, k, 1, dtype=torch.float64), |
| _t(mem["mem_ts"]).double()], 1), |
| "post_memory_egopose": post_ego, |
| "post_memory_velo": torch.cat([rec_velo, _t(mem["mem_velo"])], 1), |
| "topk_idx": idx, "rec_score": rec_score} |
|
|
| def __call__(self, x: Any, pos_embed: Any, cone: Any, E: Any, Einv: Any, pre: Dict[str, Any], |
| taps: TapRegistry = NULL_TAPS, topk_idx: Any = None) -> Dict[str, Any]: |
| """The head graph: ``pre`` holds the five ``pre_memory_*`` inputs (ONNX names). Returns the 14 ONNX outputs |
| (``post_memory_timestamp`` float64 as in the graph) + ``topk_idx`` / ``rec_score``.""" |
| mem_ego, mem_ref = self.memory_pre_update(pre, Einv, taps) |
| memory, pos = self.tokens(x, pos_embed, cone, taps) |
| t = self.temporal(mem_ref, mem_ego, pre["pre_memory_timestamp"], pre["pre_memory_embedding"], |
| pre["pre_memory_velo"], taps) |
| q = self.queries(t, taps) |
| outs_dec = self.decoder(q["tgt"], q["query_pos"], q["temp_memory"], q["temp_pos"], memory, pos, taps) |
| cls, reg = self.predict(outs_dec, q["reference_points"], taps) |
| mem = {"mem_ego": mem_ego, "mem_ref": mem_ref, "mem_emb": pre["pre_memory_embedding"], |
| "mem_ts": pre["pre_memory_timestamp"], "mem_velo": pre["pre_memory_velo"]} |
| post = self.post_update(cls, reg, outs_dec, mem, E, taps, topk_idx=topk_idx) |
| out = {"all_cls_scores": cls.flatten(0, 2).unsqueeze(0).transpose(2, 1), |
| "all_bbox_preds": reg.flatten(0, 2).unsqueeze(0).transpose(2, 1), |
| "reference_points": q["reference_points"], "tgt": q["tgt"], "temp_memory": q["temp_memory"], |
| "temp_pos": q["temp_pos"], "query_pos": q["query_pos"], "query_pos_in": self.query_pos_in, |
| "outs_dec": outs_dec} |
| out.update(post) |
| taps.tap("all_cls_scores", out["all_cls_scores"]) |
| taps.tap("all_bbox_preds", out["all_bbox_preds"]) |
| return out |
|
|
|
|
| |
|
|
| class StreamPETRNet: |
| """The three graphs over one set of weights (``weights_dir`` = the HF camera_streampetr v1.0 files).""" |
|
|
| def __init__(self, weights_dir: Any, cfg: Optional[StreamPETRConfig] = None, *, threads: Optional[int] = None, |
| check: bool = True): |
| torch = _torch() |
| if threads: |
| torch.set_num_threads(int(threads)) |
| self.cfg = cfg or default_config() |
| self.weights = StreamPETRWeights(weights_dir, check=check) |
| self.backbone = Backbone(self.weights, self.cfg) |
| self.pe = PositionEmbedding(self.weights, self.cfg) |
| self.head = Head(self.weights, self.cfg) |
|
|
| @staticmethod |
| def no_grad(): |
| return _torch().no_grad() |
|
|
