| |
| """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.<name>.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) |
| 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)) |
|
|
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
