# SPDX-License-Identifier: Apache-2.0 """BEVFormer-small in torch fp32: the deployed ``forward_trt`` graph (DerryHub/BEVFormer_tensorrt, the export behind Autoware's ``bevformer_small.onnx``), built from :class:`~.weights.BEVFormerWeights` with the backbone BatchNorms folded exactly where the port folds them. Ported from the verified research reference (``research/bevformer/scripts/bevformer_ref.py``; research/bevformer/SPEC.md sections 4 and 6). Two entry points with the same outputs ``(bev_embed [22500, 1, 256], outputs_classes [6, 1, 900, 10], outputs_coords [6, 1, 900, 10])``: - :meth:`BEVFormerNet.forward` takes the five ONNX inputs (``image``, ``prev_bev``, ``use_prev_bev``, ``can_bus``, ``lidar2img``) and computes the geometry INSIDE the graph in fp32, as the deployed graph does: the oracle form, compared with ONNX Runtime on the derived pure ONNX (``tests/test_reference_cpu.py``). - :meth:`BEVFormerNet.forward_port` takes the geometry the TT port computes on the HOST (``host.geometry``: ``ref_cam`` / SCA tables from ``lidar2img`` in float64, the TSA shift, the prev_bev rotation index map) and runs the port's exact rewrites: the rotation as a row gather with a zero row, the spatial cross-attention REBATCHED per camera into ``SCA_CAPACITY`` rows and scatter-added back by a gather-form table (no dense 6 x 22,500 pass). Proven equal to :meth:`forward` on CPU before any device work (``tests/test_reference_host.py``). Per-module taps (``self.taps = {}`` before a call; ``None`` disables them): ``backbone.stem``, ``backbone.layer1-4``, ``neck``, ``enc.can_bus_mlp``, ``enc.bev_queries``, ``enc.value``, ``enc.prev_bev_rot``, ``enc.shift``, ``ref_cam``, ``bev_mask_n``, ``encoder.layer{i}.{tsa,norm0,sca,norm1,ffn}``, ``encoder.layer{i}`` (layer output), ``decoder.layer{i}.{self_attn,norm0,cross_attn,norm1,ffn,ref}``, ``decoder.layer{i}``, ``head.cls{i}``, ``head.reg{i}``. The research goldens use the same names for the ones they hold. """ from __future__ import annotations import math from dataclasses import dataclass from typing import Any, Dict, Optional, Tuple import numpy as np import torch import torch.nn.functional as F from . import config as C from . import ops from .weights import BEVFormerWeights __all__ = ["BEVFormerNet", "PortGeometry", "reference_points_3d", "ref_2d", "point_sampling", "can_bus_shift", "inverse_sigmoid_mmdet", "inverse_sigmoid_dec"] E = C.EMBED Q = C.NUM_QUERY_BEV # ------------------------------------------------------------------------------------------------ geometry (fp32) def reference_points_3d() -> torch.Tensor: """``[1, 4, Q, 3]`` normalised (x, y, z) of the 4 pillar points of every BEV cell (DH encoder.py:169-195).""" Z = C.PC_RANGE[5] - C.PC_RANGE[2] zs = torch.linspace(0.5, Z - 0.5, C.PILLAR_POINTS).view(-1, 1, 1).expand(C.PILLAR_POINTS, C.BEV_H, C.BEV_W) / Z xs = torch.linspace(0.5, C.BEV_W - 0.5, C.BEV_W).view(1, 1, C.BEV_W).expand(C.PILLAR_POINTS, C.BEV_H, C.BEV_W) ys = torch.linspace(0.5, C.BEV_H - 0.5, C.BEV_H).view(1, C.BEV_H, 1).expand(C.PILLAR_POINTS, C.BEV_H, C.BEV_W) return torch.stack((xs / C.BEV_W, ys / C.BEV_H, zs), -1).reshape(1, C.PILLAR_POINTS, -1, 3) def ref_2d() -> torch.Tensor: """``[1, Q, 1, 2]`` normalised BEV cell centres (x, y), row-major over (y, x).""" return reference_points_3d()[0, 0, :, :2].reshape(1, -1, 1, 2) def point_sampling(ref_3d: torch.Tensor, lidar2img: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """``point_sampling_trt`` (DH encoder.py:197-259) in fp32 -> ``ref_cam [6, Q, 4, 2]``, ``bev_mask_n [6, Q, 1]`` (the any-of-4 camera mask normalised by the number of cameras that see the cell, clamp 1e-4).""" pr = C.PC_RANGE rp = ref_3d * torch.tensor([pr[3] - pr[0], pr[4] - pr[1], pr[5] - pr[2]]).view(1, 1, 1, 3) + torch.tensor(pr[:3]) rp = torch.cat((rp, torch.ones_like(rp[..., :1])), -1).view(C.PILLAR_POINTS, 1, 1, -1, 4, 1) l2i = lidar2img.view(1, 1, C.NUM_CAMS, 1, 4, 4) cam = torch.matmul(l2i, rp).squeeze(-1) # [4, 1, 6, Q, 4] eps = C.MSDA_EPS mask = (cam[..., 2:3] > eps).float() uv = cam[..., 0:2] / torch.maximum(cam[..., 2:3], torch.ones_like(cam[..., 2:3]) * eps) uv = uv / torch.tensor([C.IMG_HW[1], C.IMG_HW[0]], dtype=uv.dtype) mask = mask * (uv[..., 1:2] > 0.0) * (uv[..., 1:2] < 1.0) * (uv[..., 0:1] < 1.0) * (uv[..., 0:1] > 0.0) ref_cam = uv.permute(2, 1, 3, 0, 4).reshape(C.NUM_CAMS, -1, C.PILLAR_POINTS, 2) m = (1 - (1 - mask).prod(0)).view(C.NUM_CAMS, -1, 1) return ref_cam, m / torch.clamp(m.sum(0, keepdim=True), min=1e-4) def can_bus_shift(can_bus: torch.Tensor) -> torch.Tensor: """TSA reference shift ``[1, 2]`` (x, y) in BEV units from can_bus[0], [1], [16], with DerryHub's atan2 surrogate (``get_bev_features_trt``, DH transformer.py:262-294), in fp32 as the deployed graph.""" dx, dy = can_bus[0:1], can_bus[1:2] ego_angle = can_bus[16:17] / math.pi * 180 tl = torch.sqrt(dx ** 2 + dy ** 2) ta = (torch.atan(dy / (dx + 1e-8)) + ((1 - torch.sign(dx)) / 2) * torch.sign(dy) * math.pi) / math.pi * 180 ba = ego_angle - ta sy = tl * torch.cos(ba / 180 * math.pi) / C.GRID_LENGTH / C.BEV_H sx = tl * torch.sin(ba / 180 * math.pi) / C.GRID_LENGTH / C.BEV_W return torch.cat([sx, sy]).unsqueeze(0) def inverse_sigmoid_mmdet(x: torch.Tensor, eps: float = C.INV_SIG_EPS) -> torch.Tensor: """mmdet ``inverse_sigmoid`` (the head's coordinate decode).""" x = x.clamp(min=0, max=1) return torch.log(x.clamp(min=eps) / (1 - x).clamp(min=eps)) def inverse_sigmoid_dec(x: torch.Tensor, eps: float = C.INV_SIG_EPS) -> torch.Tensor: """The decoder's reference update (DH decoder.py:24-37).""" x = x.clamp(min=eps, max=1 - eps) return torch.log(x / (1 - x)) @dataclass class PortGeometry: """The per-frame host geometry of the TT port (``host.geometry`` builds it; all torch, fp32 / int64): ``ref_cam [6, Q, 4, 2]``; ``sca_rows [6, CAP]`` (BEV query index of every rebatched row, ``Q`` = padding); ``sca_counts [6]``; ``sca_gather [Q, K]`` (row of the stacked ``[6 * CAP + 1]`` SCA output table for each of the up to K cameras that see a query, ``6 * CAP`` = the zero row); ``sca_scale [Q, 1]`` (1 / cameras seeing it, 0 if none); ``shift [1, 2]``; ``rot_index [Q]`` (source row of the rotated prev_bev, ``Q`` = the zero row).""" ref_cam: torch.Tensor sca_rows: torch.Tensor sca_counts: torch.Tensor sca_gather: torch.Tensor sca_scale: torch.Tensor shift: torch.Tensor rot_index: torch.Tensor @classmethod def from_numpy(cls, d: Dict[str, np.ndarray]) -> "PortGeometry": def t(name, dtype): return torch.from_numpy(np.ascontiguousarray(d[name])).to(dtype) return cls(t("ref_cam", torch.float32), t("sca_rows", torch.int64), t("sca_counts", torch.int64), t("sca_gather", torch.int64), t("sca_scale", torch.float32), t("shift", torch.float32), t("rot_index", torch.int64)) # ------------------------------------------------------------------------------------------------------- model class BEVFormerNet: """The network as torch fp32 functions over the checkpoint tensors (BN folded in the backbone).""" def __init__(self, weights: BEVFormerWeights, *, q_chunk: int = 4096): self.weights = weights self.q_chunk = int(q_chunk) sd = weights.sd self.p: Dict[str, torch.Tensor] = {k: torch.from_numpy(np.array(sd[k], np.float32)) for k in sd if not k.startswith("img_backbone") and sd[k].dtype != np.int64} self.conv: Dict[str, Tuple[torch.Tensor, torch.Tensor, int, int]] = {} for module, _, _, _, _ in C.backbone_conv_specs(): cp = weights.conv(module) self.conv[module] = (torch.from_numpy(cp.weight), torch.from_numpy(cp.bias), cp.stride[0], cp.padding[0]) h = "pts_bbox_head.positional_encoding" x_embed, y_embed = self.p[f"{h}.col_embed.weight"], self.p[f"{h}.row_embed.weight"] pos = torch.cat([x_embed.unsqueeze(0).expand(C.BEV_H, C.BEV_W, -1), y_embed.unsqueeze(1).expand(C.BEV_H, C.BEV_W, -1)], dim=-1) self.bev_pos = pos.reshape(Q, E).contiguous() # [Q, 256], constant self.taps: Optional[Dict[str, torch.Tensor]] = None # ---- helpers ------------------------------------------------------------------------------------------------ def _tap(self, name: str, value: torch.Tensor) -> None: if self.taps is not None: self.taps[name] = value.detach().clone() def _lin(self, name: str, x: torch.Tensor) -> torch.Tensor: return F.linear(x, self.p[f"{name}.weight"], self.p[f"{name}.bias"]) def _ln(self, name: str, x: torch.Tensor) -> torch.Tensor: return F.layer_norm(x, (E,), self.p[f"{name}.weight"], self.p[f"{name}.bias"], eps=C.LN_EPS) def _conv(self, module: str, x: torch.Tensor, relu: bool) -> torch.Tensor: w, b, s, pad = self.conv[module] y = F.conv2d(x, w, b, stride=s, padding=pad) return F.relu(y) if relu else y # ---- image branch ------------------------------------------------------------------------------------------- def _dcn(self, module: str, x: torch.Tensor) -> torch.Tensor: """DCNv2 conv2 + folded bn2 + ReLU: offset conv (3x3 -> 27: dy/dx of 9 taps as o1 | o2, mask logits), sigmoid mask, deformable sampling + GEMM with the BN-folded weight and bias.""" out = self._conv(f"{module}.conv_offset", x, False) o1, o2, m = torch.chunk(out, 3, dim=1) w, b, _, _ = self.conv[module] return F.relu(ops.mdcn(x, torch.cat((o1, o2), dim=1), torch.sigmoid(m), w, b)) def backbone(self, x: torch.Tensor) -> torch.Tensor: """ResNet-101 (caffe style, DCNv2 in stages 3-4) on ``[6, 3, 736, 1280]`` -> C5 ``[6, 2048, 23, 40]``.""" y = F.max_pool2d(self._conv("img_backbone.conv1", x, True), 3, 2, 1) self._tap("backbone.stem", y) for s, (n, dcn) in enumerate(zip(C.RESNET_LAYERS, C.RESNET_DCN)): for b in range(n): p = f"img_backbone.layer{s + 1}.{b}" t = self._conv(f"{p}.conv1", y, True) t = self._dcn(f"{p}.conv2", t) if dcn else self._conv(f"{p}.conv2", t, True) t = self._conv(f"{p}.conv3", t, False) idt = self._conv(f"{p}.downsample.0", y, False) if b == 0 else y y = F.relu(t + idt) self._tap(f"backbone.layer{s + 1}", y) return y def neck(self, c5: torch.Tensor) -> torch.Tensor: """FPN on C5 only: lateral 1x1 (2048 -> 256) + 3x3 (256 -> 256), biases, no activation.""" n = "img_neck" lat = F.conv2d(c5, self.p[f"{n}.lateral_convs.0.conv.weight"], self.p[f"{n}.lateral_convs.0.conv.bias"]) out = F.conv2d(lat, self.p[f"{n}.fpn_convs.0.conv.weight"], self.p[f"{n}.fpn_convs.0.conv.bias"], padding=1) self._tap("neck", out) return out def extract_feat(self, image: torch.Tensor) -> torch.Tensor: return self.neck(self.backbone(image.reshape(C.NUM_CAMS, 3, C.IMG_HW[0], C.IMG_HW[1]))) # ---- encoder pieces ----------------------------------------------------------------------------------------- def bev_queries(self, can_bus: torch.Tensor) -> torch.Tensor: """``bev_embedding + can_bus_mlp(can_bus)`` -> ``[Q, 256]`` (Linear 18 -> 128, ReLU, Linear -> 256, ReLU, LayerNorm; one vector broadcast to every query).""" t = "pts_bbox_head.transformer.can_bus_mlp" cb = F.relu(self._lin(f"{t}.2", F.relu(self._lin(f"{t}.0", can_bus.view(1, 18))))) cb = self._ln(f"{t}.norm", cb) self._tap("enc.can_bus_mlp", cb.view(E)) q = self.p["pts_bbox_head.bev_embedding.weight"] + cb self._tap("enc.bev_queries", q) return q def encoder_value(self, feat: torch.Tensor) -> torch.Tensor: """``[6, 256, 23, 40]`` -> SCA value ``[6, 920, 256]`` (+ cams_embeds[cam] + level_embeds[0]).""" t = "pts_bbox_head.transformer" n, c, fh, fw = feat.shape f = feat.reshape(n, c, fh * fw).permute(0, 2, 1) v = f + self.p[f"{t}.cams_embeds"].view(n, 1, E) + self.p[f"{t}.level_embeds"][0].view(1, 1, E) self._tap("enc.value", v) return v def tsa(self, i: int, query: torch.Tensor, value2: torch.Tensor, hybrid_ref: torch.Tensor) -> torch.Tensor: """TemporalSelfAttention (+ residual): ``query [Q, 256]``, ``value2 [2, Q, 256]`` (queue: [prev; current]), ``hybrid_ref [2, Q, 1, 2]``.""" a = f"pts_bbox_head.transformer.encoder.layers.{i}.attentions.0" q = query + self.bev_pos qcat = torch.cat([value2[0], q], -1).view(1, Q, 2 * E) # [1, Q, 512] v = self._lin(f"{a}.value_proj", value2).view(C.TSA_QUEUE, Q, C.NUM_HEADS, C.HEAD_DIM) off = self._lin(f"{a}.sampling_offsets", qcat).view(1, Q, C.NUM_HEADS, C.TSA_QUEUE, 1, C.TSA_POINTS, 2) off = off.permute(0, 3, 1, 2, 4, 5, 6).reshape(C.TSA_QUEUE, Q, C.NUM_HEADS, C.TSA_POINTS * 2) aw = self._lin(f"{a}.attention_weights", qcat).view(1, Q, C.NUM_HEADS, C.TSA_QUEUE, 1, C.TSA_POINTS) aw = aw.permute(0, 3, 1, 2, 4, 5).reshape(C.TSA_QUEUE, Q, C.NUM_HEADS, C.TSA_POINTS) out = ops.msda(v, [(C.BEV_H, C.BEV_W)], hybrid_ref, off, aw, self.q_chunk).flatten(2) # [2, Q, 256] return self._lin(f"{a}.output_proj", out.mean(0)) + query def _sca_proj(self, i: int, value: torch.Tensor, q_rows: torch.Tensor): a = f"pts_bbox_head.transformer.encoder.layers.{i}.attentions.1.deformable_attention" v = self._lin(f"{a}.value_proj", value).view(C.NUM_CAMS, C.FEAT_LEN, C.NUM_HEADS, C.HEAD_DIM) off = self._lin(f"{a}.sampling_offsets", q_rows) aw = self._lin(f"{a}.attention_weights", q_rows) return v, off, aw def sca_dense(self, i: int, query: torch.Tensor, value: torch.Tensor, ref_cam: torch.Tensor, mask_n: torch.Tensor) -> torch.Tensor: """SpatialCrossAttention as the TRT graph runs it: dense over all 6 x 22,500 (camera, query) pairs, then the normalised-mask sum over cameras (+ output_proj + residual).""" qr = query.unsqueeze(0).expand(C.NUM_CAMS, Q, E) v, off, aw = self._sca_proj(i, value, qr) ref = ref_cam.reshape(C.NUM_CAMS, Q, 1, C.PILLAR_POINTS * 2) out = ops.msda(v, [C.FEAT_HW], ref, off.view(C.NUM_CAMS, Q, C.NUM_HEADS, C.SCA_POINTS * 2), aw.view(C.NUM_CAMS, Q, C.NUM_HEADS, C.SCA_POINTS), self.q_chunk).flatten(2) slots = (out * mask_n).sum(0) a = f"pts_bbox_head.transformer.encoder.layers.{i}.attentions.1" return self._lin(f"{a}.output_proj", slots) + query def sca_rebatched(self, i: int, query: torch.Tensor, value: torch.Tensor, geo: PortGeometry) -> torch.Tensor: """The port's SCA: per camera, gather its ``SCA_CAPACITY`` rows (valid queries + padding), run the linears and the MSDA on those rows only, stack the six outputs with a zero row, and sum every query's rows through the gather table, times 1 / cameras (identical to :meth:`sca_dense`: a query seen by no camera gets 0).""" cap = geo.sca_rows.shape[1] q_pad = torch.cat([query, torch.zeros(1, E, dtype=query.dtype)], 0) # row Q = padding rows = q_pad[geo.sca_rows.reshape(-1)].view(C.NUM_CAMS, cap, E) v, off, aw = self._sca_proj(i, value, rows) ref_pad = torch.cat([geo.ref_cam, torch.zeros(C.NUM_CAMS, 1, C.PILLAR_POINTS, 2)], 1) ref = torch.stack([ref_pad[c, geo.sca_rows[c]] for c in range(C.NUM_CAMS)]) # [6, cap, 4, 2] out = ops.msda(v, [C.FEAT_HW], ref.reshape(C.NUM_CAMS, cap, 1, C.PILLAR_POINTS * 2), off.view(C.NUM_CAMS, cap, C.NUM_HEADS, C.SCA_POINTS * 2), aw.view(C.NUM_CAMS, cap, C.NUM_HEADS, C.SCA_POINTS), self.q_chunk).flatten(2) table = torch.cat([out.reshape(C.NUM_CAMS * cap, E), torch.zeros(1, E, dtype=out.dtype)], 0) slots = table[geo.sca_gather].sum(1) * geo.sca_scale # [Q, 256] a = f"pts_bbox_head.transformer.encoder.layers.{i}.attentions.1" return self._lin(f"{a}.output_proj", slots) + query def ffn(self, prefix: str, x: torch.Tensor) -> torch.Tensor: return x + self._lin(f"{prefix}.ffns.0.layers.1", F.relu(self._lin(f"{prefix}.ffns.0.layers.0.0", x))) def encoder(self, bev_q: torch.Tensor, value: torch.Tensor, prev_rot: torch.Tensor, use_prev_bev: torch.Tensor, shift: torch.Tensor, sca) -> torch.Tensor: """3 x (TSA -> LN -> SCA -> LN -> FFN -> LN). ``sca(i, query) -> query`` is the dense or rebatched SCA. The TSA value is ``use * [rot(prev_bev); bev_queries] + (1 - use) * [query; query]`` with the CURRENT layer's query in the second term (DH encoder.py:305, 587-589).""" use = use_prev_bev.reshape(()) r2 = ref_2d() hybrid = torch.cat([r2 + shift.view(1, 1, 1, 2) * use, r2], 0) # [2, Q, 1, 2] prev = torch.stack([prev_rot, bev_q]) # [2, Q, 256] query = bev_q for i in range(C.ENC_LAYERS): pre = f"pts_bbox_head.transformer.encoder.layers.{i}" value2 = use * prev + (1 - use) * query.unsqueeze(0).expand(2, Q, E) query = self.tsa(i, query, value2, hybrid) self._tap(f"encoder.layer{i}.tsa", query) query = self._ln(f"{pre}.norms.0", query) self._tap(f"encoder.layer{i}.norm0", query) query = sca(i, query) self._tap(f"encoder.layer{i}.sca", query) query = self._ln(f"{pre}.norms.1", query) self._tap(f"encoder.layer{i}.norm1", query) query = self.ffn(pre, query) self._tap(f"encoder.layer{i}.ffn", query) query = self._ln(f"{pre}.norms.2", query) self._tap(f"encoder.layer{i}", query) return query # ---- decoder + heads ---------------------------------------------------------------------------------------- def _mha(self, i: int, x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: """mmcv MultiheadAttention (q = k = x + pos, v = x, 8 heads, batch_first False) + residual.""" a = f"pts_bbox_head.transformer.decoder.layers.{i}.attentions.0.attn" wq, wk, wv = self.p[f"{a}.in_proj_weight"].chunk(3, 0) bq, bk, bv = self.p[f"{a}.in_proj_bias"].chunk(3, 0) qk = x + pos n = x.shape[0] q = F.linear(qk, wq, bq).view(n, C.NUM_HEADS, C.HEAD_DIM).transpose(0, 1) k = F.linear(qk, wk, bk).view(n, C.NUM_HEADS, C.HEAD_DIM).transpose(0, 1) v = F.linear(x, wv, bv).view(n, C.NUM_HEADS, C.HEAD_DIM).transpose(0, 1) att = torch.softmax(torch.matmul(q, k.transpose(1, 2)) * (1.0 / math.sqrt(C.HEAD_DIM)), dim=-1) o = torch.matmul(att, v).transpose(0, 1).reshape(n, E) return x + self._lin(f"{a}.out_proj", o) def _cross(self, i: int, x: torch.Tensor, pos: torch.Tensor, bev: torch.Tensor, ref_xy: torch.Tensor): """CustomMSDeformableAttention over the BEV (150 x 150, 4 points at ``ref_xy``) + residual.""" a = f"pts_bbox_head.transformer.decoder.layers.{i}.attentions.1" q = (x + pos).view(1, -1, E) n = q.shape[1] v = self._lin(f"{a}.value_proj", bev.view(1, Q, E)).view(1, Q, C.NUM_HEADS, C.HEAD_DIM) off = self._lin(f"{a}.sampling_offsets", q).view(1, n, C.NUM_HEADS, C.DEC_POINTS * 2) aw = self._lin(f"{a}.attention_weights", q).view(1, n, C.NUM_HEADS, C.DEC_POINTS) out = ops.msda(v, [(C.BEV_H, C.BEV_W)], ref_xy.view(1, n, 1, 2), off, aw, self.q_chunk).flatten(2) return x + self._lin(f"{a}.output_proj", out.view(n, E)) def _cls(self, i: int, h: torch.Tensor) -> torch.Tensor: b = f"pts_bbox_head.cls_branches.{i}" h = F.relu(self._ln(f"{b}.1", self._lin(f"{b}.0", h))) h = F.relu(self._ln(f"{b}.4", self._lin(f"{b}.3", h))) return self._lin(f"{b}.6", h) def _reg(self, i: int, h: torch.Tensor) -> torch.Tensor: b = f"pts_bbox_head.reg_branches.{i}" h = F.relu(self._lin(f"{b}.2", F.relu(self._lin(f"{b}.0", h)))) return self._lin(f"{b}.4", h) def decode(self, bev: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """6 decoder layers (MHA -> LN -> deformable cross-attention -> LN -> FFN -> LN, box refinement) + the six cls / reg heads -> ``outputs_classes``, ``outputs_coords`` ``[6, 1, 900, 10]`` (cx, cy, cz in metres).""" qe = self.p["pts_bbox_head.query_embedding.weight"] pos, x = qe[:, :E], qe[:, E:] ref = torch.sigmoid(self._lin("pts_bbox_head.transformer.reference_points", pos)) # [900, 3] init_ref = ref cls_all, crd_all = [], [] pr = C.PC_RANGE for i in range(C.DEC_LAYERS): pre = f"pts_bbox_head.transformer.decoder.layers.{i}" ref_in = ref x = self._mha(i, x, pos) self._tap(f"decoder.layer{i}.self_attn", x) x = self._ln(f"{pre}.norms.0", x) self._tap(f"decoder.layer{i}.norm0", x) x = self._cross(i, x, pos, bev, ref[:, :2]) self._tap(f"decoder.layer{i}.cross_attn", x) x = self._ln(f"{pre}.norms.1", x) self._tap(f"decoder.layer{i}.norm1", x) x = self.ffn(pre, x) self._tap(f"decoder.layer{i}.ffn", x) x = self._ln(f"{pre}.norms.2", x) self._tap(f"decoder.layer{i}", x) t = self._reg(i, x) # the head's reg == the update's ref = torch.sigmoid(torch.cat([t[:, 0:2] + inverse_sigmoid_dec(ref[:, :2]), t[:, 4:5] + inverse_sigmoid_dec(ref[:, 2:3])], -1)) self._tap(f"decoder.layer{i}.ref", ref) r = inverse_sigmoid_mmdet(init_ref if i == 0 else ref_in) cls = self._cls(i, x) xy = torch.sigmoid(t[:, 0:2] + r[:, 0:2]) z = torch.sigmoid(t[:, 4:5] + r[:, 2:3]) crd = torch.cat([xy[:, 0:1] * (pr[3] - pr[0]) + pr[0], xy[:, 1:2] * (pr[4] - pr[1]) + pr[1], t[:, 2:4], z * (pr[5] - pr[2]) + pr[2], t[:, 5:]], -1) self._tap(f"head.cls{i}", cls) self._tap(f"head.reg{i}", t) cls_all.append(cls) crd_all.append(crd) return torch.stack(cls_all).unsqueeze(1), torch.stack(crd_all).unsqueeze(1) # ---- entry points ------------------------------------------------------------------------------------------- @torch.no_grad() def forward(self, image, prev_bev, use_prev_bev, can_bus, lidar2img): """The deployed graph: the five ONNX inputs (any array-likes, float32) -> the three ONNX outputs (torch).""" image, prev_bev, use_prev_bev, can_bus, lidar2img = (torch.as_tensor(np.asarray(a, np.float32)) for a in (image, prev_bev, use_prev_bev, can_bus, lidar2img)) feat = self.extract_feat(image) value = self.encoder_value(feat) bev_q = self.bev_queries(can_bus) shift = can_bus_shift(can_bus) self._tap("enc.shift", shift.view(2)) pb = prev_bev.reshape(C.BEV_H, C.BEV_W, E).permute(2, 0, 1) pb = ops.rotate(pb, can_bus[17:18], C.ROTATE_CENTER).permute(1, 2, 0).reshape(Q, E) self._tap("enc.prev_bev_rot", pb) ref_cam, mask_n = point_sampling(reference_points_3d(), lidar2img) self._tap("ref_cam", ref_cam) self._tap("bev_mask_n", mask_n) bev = self.encoder(bev_q, value, pb, use_prev_bev, shift, lambda i, q: self.sca_dense(i, q, value, ref_cam, mask_n)) cls, crd = self.decode(bev) return bev.view(Q, 1, E), cls, crd @torch.no_grad() def forward_port(self, image, prev_bev, use_prev_bev, can_bus, geo: PortGeometry): """The TT port's decomposition with host geometry ``geo`` (see :class:`PortGeometry`).""" image, prev_bev, use_prev_bev, can_bus = (torch.as_tensor(np.asarray(a, np.float32)) for a in (image, prev_bev, use_prev_bev, can_bus)) feat = self.extract_feat(image) value = self.encoder_value(feat) bev_q = self.bev_queries(can_bus) self._tap("enc.shift", geo.shift.view(2)) table = torch.cat([prev_bev.reshape(Q, E), torch.zeros(1, E)], 0) pb = table[geo.rot_index] # row gather, zero row self._tap("enc.prev_bev_rot", pb) self._tap("ref_cam", geo.ref_cam) bev = self.encoder(bev_q, value, pb, use_prev_bev, geo.shift, lambda i, q: self.sca_rebatched(i, q, value, geo)) cls, crd = self.decode(bev) return bev.view(Q, 1, E), cls, crd def run(self, *args: Any, taps: bool = False, port: bool = False): """``forward`` / ``forward_port`` returning numpy plus the taps (numpy) when ``taps``.""" self.taps = {} if taps else None try: outs = (self.forward_port if port else self.forward)(*args) tp = {k: v.numpy() for k, v in (self.taps or {}).items()} finally: self.taps = None return tuple(o.numpy() for o in outs), tp