# SPDX-License-Identifier: Apache-2.0 """Pure-PyTorch fp32 re-implementation of the plugin-free parts of ``bevdet_one_lt_d.onnx`` with per-module taps. Ported from the verified research reference (``research/bevdet/scripts/bevdet_torch.py``: PCC 1.0 against ONNX Runtime on the deployed sub-graphs, SPEC section 6). Module boundaries mirror the training code (LCH1238/BEVDet@export: ResNet-50, CustomFPN, DepthNet (BEVDepth), CustomResNet, FPN_LSS, CenterHead) so that each TT stage has a reference tap. Every convolution runs with the weights and the attributes of its ONNX node (``reference.weights``). ``BEVDetNet.image_encoder(img, cam_params)``: normalised images [6, 3, 256, 704] + ``cam_params`` [1, 6, 27] -> (depth [6, 118, 16, 44] after the softmax over depth bins, feat [6, 80, 16, 44]); ``gates=`` replaces the camera MLP / SE branch by precomputed per-camera gates (the TT port's form, ``reference.port_form.camera_gates``). ``BEVDetNet.bev_decoder(bev_in)``: GatherBEV output [1, 720, 128, 128] -> the 6 head maps [1, k, 128, 128]. Tap names (``ttaw.golden.TapRegistry``; the device PCC gates use the same names): ============================= ===================== ============================================================ tap shape where ============================= ===================== ============================================================ ``img.backbone.stem`` [6, 64, 64, 176] conv1 7x7 s2 + ReLU + MaxPool 3x3 s2 ``img.backbone.layer1..4`` [6, 256..2048, ...] ResNet-50 stages (layer3 = C4 [6,1024,16,44], layer4 = C5) ``img.neck`` [6, 512, 16, 44] CustomFPN output (fpn conv, no activation) ``depthnet.reduce`` [6, 512, 16, 44] reduce_conv 3x3 + ReLU ``depthnet.gate_ctx/_depth`` [6, 512] sigmoid(SE(MLP(BN1d(cam_params)))), calibration-only ``feat`` [6, 80, 16, 44] context_conv(reduce * gate_ctx): the BEVPool features ``depthnet.blocks.0..2`` [6, 512, 16, 44] BasicBlocks on reduce * gate_depth ``depthnet.aspp.b1..b4/.gap`` [6, 96, 16, 44] ASPP branches (dilation 1/6/12/18, GAP: [6, 96, 1, 1]) ``depthnet.aspp`` [6, 512, 16, 44] ASPP conv1 (480 -> 512) + ReLU ``depthnet.logits`` [6, 118, 16, 44] depth_conv.4 1x1 ``depth`` [6, 118, 16, 44] softmax over the 118 bins ``bev.backbone.0..2`` [1, 160, 64, 64] ... CustomResNet stages (F0, F1, F2 [1, 640, 16, 16]) ``bev.neck.up4`` [1, 640, 64, 64] bilinear x4 (align_corners) of F2 ``bev.neck.conv`` [1, 512, 64, 64] two 3x3 + ReLU on cat(F0, up4) ``bev.neck.up2`` [1, 512, 128, 128] bilinear x2 (align_corners) ``bev.neck`` [1, 256, 128, 128] 3x3 + ReLU, 1x1 ``head.shared`` [1, 64, 128, 128] shared 3x3 + ReLU ``head..hidden`` [1, 64, 128, 128] per-branch 3x3 + ReLU ``head.reg_0`` ... ``heatmap_0`` [1, k, 128, 128] the ONNX outputs (``HEAD_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 HEAD_NAMES, BEVDetConfig, default_config from .weights import HEAD_BRANCHES, BEVDetWeights __all__ = ["BEVDetNet"] def _torch(): import torch return torch class BEVDetNet: """The network as fp32 torch functions over the ONNX weights (no autograd, no nn.Module state).""" def __init__(self, weights: BEVDetWeights, cfg: Optional[BEVDetConfig] = None, *, dtype: Any = None): torch = _torch() self.cfg = cfg or default_config() self.cfg.check() self.weights = weights self.dtype = dtype or torch.float32 self._params: Dict[str, Tuple[Any, Optional[Any], Any]] = {} for module, conv in weights.iter_convs(): w = torch.from_numpy(np.array(conv.weight, np.float32)).to(self.dtype) # own copy (writable) b = None if conv.bias is None else torch.from_numpy(np.array(conv.bias, np.float32)).to(self.dtype) self._params[module] = (w, b, conv) dn = "img_view_transformer.depth_net." self._lin = {m: tuple(torch.tensor(np.array(a, np.float32), dtype=self.dtype) for a in weights.linear(dn + m)) for m in ("context_mlp.fc1", "context_mlp.fc2", "depth_mlp.fc1", "depth_mlp.fc2")} bn = weights.depth_bn() self._bn = tuple(torch.from_numpy(np.asarray(a, np.float32)).to(self.dtype) for a in (bn.gamma, bn.beta, bn.mean, bn.var)) self._bn_eps = float(np.float32(bn.eps)) # ---- primitives ----------------------------------------------------------------------------------- def conv(self, x: Any, module: str, relu: bool = False) -> Any: torch = _torch() import torch.nn.functional as F w, b, c = self._params[module] y = F.conv2d(x, w, b, stride=c.stride, padding=c.padding, dilation=c.dilation, groups=c.groups) return torch.relu(y) if relu else y def linear(self, x: Any, module: str, relu: bool = False) -> Any: torch = _torch() w, b = self._lin[module] y = torch.nn.functional.linear(x, w, b) return torch.relu(y) if relu else y # ---- image encoder -------------------------------------------------------------------------------- def backbone(self, x: Any, taps: TapRegistry = NULL_TAPS) -> Tuple[Any, Any]: """ResNet-50 (pytorch style: stride on the 3x3, 1x1 downsample, BN folded) -> (C4, C5).""" torch = _torch() import torch.nn.functional as F x = self.conv(x, "img_backbone.conv1", relu=True) x = taps.tap("img.backbone.stem", F.max_pool2d(x, 3, 2, 1)) outs = [] for li, nb in enumerate([3, 4, 6, 3], start=1): for bi in range(nb): p = f"img_backbone.layer{li}.{bi}" idt = self.conv(x, p + ".downsample.0") if bi == 0 else x y = self.conv(x, p + ".conv1", relu=True) y = self.conv(y, p + ".conv2", relu=True) y = self.conv(y, p + ".conv3") x = torch.relu(y + idt) outs.append(taps.tap(f"img.backbone.layer{li}", x)) return outs[2], outs[3] def neck(self, c4: Any, c5: Any, taps: TapRegistry = NULL_TAPS) -> Any: """CustomFPN: L0 = 1x1(C4) + nearest-upsample(1x1(C5)) (ONNX Resize nearest / asymmetric / floor), then the 3x3 fpn conv with no activation.""" import torch.nn.functional as F l0 = self.conv(c4, "img_neck.lateral_convs.0.conv") l1 = self.conv(c5, "img_neck.lateral_convs.1.conv") l0 = l0 + F.interpolate(l1, size=l0.shape[2:], mode="nearest") return taps.tap("img.neck", self.conv(l0, "img_neck.fpn_convs.0.conv")) def camera_mlp_input(self, cam_params: Any) -> Any: """``depth_net.bn`` (BatchNorm1d(27), float32 like the graph) on ``cam_params.reshape(-1, 27)``.""" torch = _torch() g, b, m, v = self._bn x = cam_params.reshape(-1, cam_params.shape[-1]).to(self.dtype) return (x - m) / torch.sqrt(v + self._bn_eps) * g + b def camera_gates(self, cam_params: Any) -> Tuple[Any, Any]: """The deployed graph's camera branch: ``sigmoid(SE(MLP(BN1d(cam_params))))`` for context and depth -> ([6, 512], [6, 512]) float32.""" torch = _torch() dn = "img_view_transformer.depth_net." mlp_in = self.camera_mlp_input(cam_params) def gate(mlp: str, se: str) -> Any: v = self.linear(self.linear(mlp_in, mlp + ".fc1", relu=True), mlp + ".fc2")[..., None, None] g = self.conv(self.conv(v, dn + se + ".conv_reduce", relu=True), dn + se + ".conv_expand") return torch.sigmoid(g)[..., 0, 0] return gate("context_mlp", "context_se"), gate("depth_mlp", "depth_se") def depth_net(self, x: Any, cam_params: Any = None, *, gates: Optional[Tuple[Any, Any]] = None, taps: TapRegistry = NULL_TAPS) -> Tuple[Any, Any]: """DepthNet (BEVDepth, camera-aware): -> (depth softmax [6, 118, 16, 44], feat [6, 80, 16, 44]).""" torch = _torch() import torch.nn.functional as F dn = "img_view_transformer.depth_net." x = taps.tap("depthnet.reduce", self.conv(x, dn + "reduce_conv.0", relu=True)) if gates is None: gates = self.camera_gates(cam_params) g_ctx = taps.tap("depthnet.gate_ctx", gates[0]) g_depth = taps.tap("depthnet.gate_depth", gates[1]) feat = taps.tap("feat", self.conv(x * g_ctx[:, :, None, None], dn + "context_conv")) d = x * g_depth[:, :, None, None] for bi in range(3): q = dn + f"depth_conv.{bi}" y = self.conv(d, q + ".conv1", relu=True) y = self.conv(y, q + ".conv2") d = taps.tap(f"depthnet.blocks.{bi}", torch.relu(y + d)) q = dn + "depth_conv.3" x1 = taps.tap("depthnet.aspp.b1", self.conv(d, q + ".aspp1.atrous_conv", relu=True)) x2 = taps.tap("depthnet.aspp.b2", self.conv(d, q + ".aspp2.atrous_conv", relu=True)) x3 = taps.tap("depthnet.aspp.b3", self.conv(d, q + ".aspp3.atrous_conv", relu=True)) x4 = taps.tap("depthnet.aspp.b4", self.conv(d, q + ".aspp4.atrous_conv", relu=True)) x5 = self.conv(F.adaptive_avg_pool2d(d, 1), q + ".global_avg_pool.1", relu=True) taps.tap("depthnet.aspp.gap", x5) x5 = F.interpolate(x5, size=x4.shape[2:], mode="bilinear", align_corners=True) d = taps.tap("depthnet.aspp", self.conv(torch.cat([x1, x2, x3, x4, x5], 1), q + ".conv1", relu=True)) logits = taps.tap("depthnet.logits", self.conv(d, dn + "depth_conv.4")) return taps.tap("depth", logits.softmax(dim=1)), feat def image_encoder(self, img: Any, cam_params: Any = None, *, gates: Optional[Tuple[Any, Any]] = None, taps: TapRegistry = NULL_TAPS) -> Tuple[Any, Any]: """[6, 3, 256, 704] normalised images (+ ``cam_params`` [1, 6, 27] or precomputed ``gates``) -> (depth [6, 118, 16, 44], feat [6, 80, 16, 44]).""" torch = _torch() with torch.no_grad(): img = torch.as_tensor(img).to(self.dtype) cp = None if cam_params is None else torch.as_tensor(cam_params).to(self.dtype) if gates is not None: gates = tuple(torch.as_tensor(g).to(self.dtype) for g in gates) c4, c5 = self.backbone(img, taps) x = self.neck(c4, c5, taps) return self.depth_net(x, cp, gates=gates, taps=taps) # ---- BEV decoder ---------------------------------------------------------------------------------- def bev_encoder(self, x: Any, taps: TapRegistry = NULL_TAPS) -> Any: """CustomResNet (BasicBlocks, 3x3 s2 downsample) + FPN_LSS (bilinear align_corners x4 / x2).""" torch = _torch() import torch.nn.functional as F feats = [] for li in range(3): for bi in range(2): p = f"img_bev_encoder_backbone.layers.{li}.{bi}" idt = self.conv(x, p + ".downsample") if bi == 0 else x y = self.conv(x, p + ".conv1", relu=True) y = self.conv(y, p + ".conv2") x = torch.relu(y + idt) feats.append(taps.tap(f"bev.backbone.{li}", x)) x2 = taps.tap("bev.neck.up4", F.interpolate(feats[2], scale_factor=4, mode="bilinear", align_corners=True)) x = torch.cat([feats[0], x2], 1) x = self.conv(x, "img_bev_encoder_neck.conv.0", relu=True) x = taps.tap("bev.neck.conv", self.conv(x, "img_bev_encoder_neck.conv.3", relu=True)) x = taps.tap("bev.neck.up2", F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)) x = self.conv(x, "img_bev_encoder_neck.up2.1", relu=True) return taps.tap("bev.neck", self.conv(x, "img_bev_encoder_neck.up2.4")) def center_head(self, x: Any, taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]: x = taps.tap("head.shared", self.conv(x, "pts_bbox_head.shared_conv.conv", relu=True)) out = {} for h, name in zip(HEAD_BRANCHES, HEAD_NAMES): p = f"pts_bbox_head.task_heads.0.{h}" y = taps.tap(f"head.{name}.hidden", self.conv(x, p + ".0.conv", relu=True)) out[name] = taps.tap(f"head.{name}", self.conv(y, p + ".1")) return out def bev_decoder(self, bev_in: Any, taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]: """[1, 720, 128, 128] -> {``reg_0``: [1, 2, 128, 128], ..., ``heatmap_0``: [1, 10, 128, 128]}.""" torch = _torch() with torch.no_grad(): x = torch.as_tensor(bev_in).to(self.dtype) return self.center_head(self.bev_encoder(x, taps), taps)