# SPDX-License-Identifier: Apache-2.0 """Weights of ``bevdet_one_lt_d.onnx`` (AutowareFoundation/tensorrt_bevdet@v1.0), read the way the TT port reads them. The ONNX is parsed as data with the vendored ``ttaw.weights.OnnxWeights`` (C04): nothing is executed, and every parameter is addressed through the graph node that consumes it (``consumer_of(, op)`` -> the node -> ``conv()`` / ``gemm()`` / ``batchnorm()``), so the convolution attributes (strides, pads, dilations, groups) come from the deployed graph itself, never from re-typed constants. The export folded every BatchNorm2d into its convolution (``fuse_conv_bn=True``, S:271), so the conv weights are used as stored; the one remaining BatchNorm (``depth_net.bn`` = BatchNorm1d(27) over ``cam_params``) is folded by the port at calibration time into the camera gates (``reference.port_form.camera_gates``, float64, one final rounding), exactly where the TT port folds it. :meth:`BEVDetWeights.check_graph` verifies that the file is the network this port was written for (node coverage, plugin attributes, resize modes, softmax axis, output names) and fails loudly otherwise. """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, Iterator, List, Optional, Tuple import numpy as np from .config import HEAD_NAMES, BEVDetConfig __all__ = ["ONNX_FILE", "ONNX_SHA256", "BEVDetWeights", "module_names", "find_onnx"] ONNX_FILE = "bevdet_one_lt_d.onnx" # sha256 of the pinned file (HF LFS sha256 == local, SPEC section 1) ONNX_SHA256 = "47f156cb63ca97fa1d47b5d46eba648d4bbe71c95a08af14d3d0fdf53d800d77" HEAD_BRANCHES = ("reg", "height", "dim", "rot", "vel", "heatmap") def module_names() -> Dict[str, List[str]]: """Every convolution / linear module of the network by stage, in graph order (initializer prefixes).""" backbone = ["img_backbone.conv1"] for li, nb in enumerate([3, 4, 6, 3], start=1): for bi in range(nb): p = f"img_backbone.layer{li}.{bi}" backbone += [f"{p}.conv1", f"{p}.conv2", f"{p}.conv3"] + ([f"{p}.downsample.0"] if bi == 0 else []) neck = ["img_neck.lateral_convs.0.conv", "img_neck.lateral_convs.1.conv", "img_neck.fpn_convs.0.conv"] dn = "img_view_transformer.depth_net." depthnet = [dn + "reduce_conv.0", dn + "context_se.conv_reduce", dn + "context_se.conv_expand", dn + "context_conv", dn + "depth_se.conv_reduce", dn + "depth_se.conv_expand"] for bi in range(3): depthnet += [dn + f"depth_conv.{bi}.conv1", dn + f"depth_conv.{bi}.conv2"] depthnet += [dn + f"depth_conv.3.aspp{i}.atrous_conv" for i in range(1, 5)] depthnet += [dn + "depth_conv.3.global_avg_pool.1", dn + "depth_conv.3.conv1", dn + "depth_conv.4"] linears = [dn + "context_mlp.fc1", dn + "context_mlp.fc2", dn + "depth_mlp.fc1", dn + "depth_mlp.fc2"] bev = [] for li in range(3): for bi in range(2): p = f"img_bev_encoder_backbone.layers.{li}.{bi}" bev += [f"{p}.conv1", f"{p}.conv2"] + ([f"{p}.downsample"] if bi == 0 else []) bev_neck = ["img_bev_encoder_neck.conv.0", "img_bev_encoder_neck.conv.3", "img_bev_encoder_neck.up2.1", "img_bev_encoder_neck.up2.4"] head = ["pts_bbox_head.shared_conv.conv"] for h in HEAD_BRANCHES: head += [f"pts_bbox_head.task_heads.0.{h}.0.conv", f"pts_bbox_head.task_heads.0.{h}.1"] return {"img_backbone": backbone, "img_neck": neck, "depth_net": depthnet, "depth_net_linear": linears, "bev_backbone": bev, "bev_neck": bev_neck, "head": head} @dataclass(frozen=True) class Conv: """One convolution as the deployed graph runs it (float32 weights as stored).""" node: str weight: np.ndarray # [Cout, Cin / groups, kh, kw] bias: Optional[np.ndarray] # [Cout] stride: Tuple[int, int] pads: Tuple[int, int, int, int] # ONNX order: top, left, bottom, right dilation: Tuple[int, int] groups: int @property def padding(self) -> Tuple[int, int]: t, l, b, r = self.pads if (t, l) != (b, r): raise ValueError(f"{self.node}: asymmetric pads {self.pads}") return t, l class BEVDetWeights: """Parameters of the deployed network, addressed by consuming node. ``conv(module)`` / ``linear(module)`` take the module prefix of the training code (e.g. ``img_backbone.layer1.0.conv1``).""" def __init__(self, onnx_path: Any, *, check: bool = True): from ..ttaw.weights import OnnxWeights self.path = Path(onnx_path) self.onnx = OnnxWeights(self.path) self._convs: Dict[str, Conv] = {} if check: self.check_graph() @property def sha256(self) -> str: return self.onnx.sha256 # ---- addressing ---------------------------------------------------------------------------------- def conv_node(self, module: str) -> str: return self.onnx.consumer_of(module + ".weight", "Conv").name def conv(self, module: str) -> Conv: if module not in self._convs: node = self.conv_node(module) p = self.onnx.conv(node) if p.auto_pad not in ("NOTSET", ""): raise ValueError(f"{node}: auto_pad {p.auto_pad} is not supported") self._convs[module] = Conv(node, np.asarray(p.weight, np.float32), None if p.bias is None else np.asarray(p.bias, np.float32), tuple(p.strides), tuple(p.pads), tuple(p.dilations), int(p.group)) return self._convs[module] def linear(self, module: str) -> Tuple[np.ndarray, np.ndarray]: """Gemm (``transB=1``, alpha = beta = 1) -> (weight [out, in], bias [out]) float32.""" node = self.onnx.consumer_of(module + ".weight", "Gemm").name g = self.onnx.gemm(node) if g.trans_a or not g.trans_b or g.alpha != 1.0 or g.beta != 1.0: raise ValueError(f"{node}: unexpected Gemm attributes {g}") return np.asarray(g.weight, np.float32), np.asarray(g.bias, np.float32) def depth_bn(self): """``depth_net.bn``: BatchNorm1d(27) over ``cam_params`` (``ttaw.weights.BatchNorm``, float64).""" node = self.onnx.consumer_of("img_view_transformer.depth_net.bn.weight", "BatchNormalization").name return self.onnx.batchnorm(node) def iter_convs(self) -> Iterator[Tuple[str, Conv]]: for modules in module_names().values(): for m in modules: if m.endswith(("fc1", "fc2")): continue yield m, self.conv(m) # ---- graph checks --------------------------------------------------------------------------------- def check_graph(self, cfg: Optional[BEVDetConfig] = None) -> None: """The file is the deployed BEVDet this port reproduces: every Conv / Gemm node maps to exactly one module, the 4 ``bevdet::`` plugins carry the expected attributes, the resize / softmax semantics are the ones the reference implements, and the outputs are the 6 head maps + ``curr_bevfeat``.""" from .config import default_config cfg = cfg or default_config() w = self.onnx names = module_names() conv_modules = [m for k, v in names.items() if k != "depth_net_linear" for m in v] used = [self.conv_node(m) for m in conv_modules] all_convs = [n.name for n in w.nodes("Conv")] if sorted(used) != sorted(all_convs): raise ValueError(f"{self.path.name}: Conv nodes do not map 1:1 to the BEVDet modules " f"({len(set(used))} mapped, {len(all_convs)} in the graph)") gemms = sorted(w.consumer_of(m + ".weight", "Gemm").name for m in names["depth_net_linear"]) if gemms != sorted(n.name for n in w.nodes("Gemm")): raise ValueError(f"{self.path.name}: Gemm nodes do not map to the 4 camera MLP layers") pre = w.find_node("*", "Preprocess") r = float(np.float32(pre.attrs["resize_radio"])) if (int(pre.attrs["crop_h"]), int(pre.attrs["crop_w"])) != cfg.crop_hw or r != float(cfg.resize): raise ValueError(f"Preprocess attributes {pre.attrs} differ from the configuration") bp = w.find_node("*", "BEVPool") ny, nx = cfg.bev_hw if (int(bp.attrs["bev_h"]), int(bp.attrs["bev_w"]), int(bp.attrs["n"])) != (ny, nx, cfg.n_cams): raise ValueError(f"BEVPool attributes {bp.attrs} differ from the configuration") w.find_node("*", "AlignBEV") w.find_node("*", "GatherBEV") modes = [(n.attrs.get("mode"), n.attrs.get("coordinate_transformation_mode"), n.attrs.get("nearest_mode")) for n in w.nodes("Resize")] expect = [("nearest", "asymmetric", "floor"), ("linear", "align_corners", "floor"), ("linear", "align_corners", "floor"), ("linear", "align_corners", "floor")] if modes != expect: raise ValueError(f"Resize modes {modes} differ from the implemented ones {expect}") sm = w.nodes("Softmax") if len(sm) != 1 or int(sm[0].attrs.get("axis", -1)) != 1: raise ValueError("expected one Softmax over axis 1 (depth bins)") mp = w.nodes("MaxPool") if len(mp) != 1 or tuple(mp[0].attrs["kernel_shape"]) != (3, 3) or tuple(mp[0].attrs["strides"]) != (2, 2) \ or tuple(mp[0].attrs["pads"]) != (1, 1, 1, 1) or int(mp[0].attrs.get("ceil_mode", 0)) != 0: raise ValueError("expected the ResNet stem MaxPool 3x3 s2 p1") if tuple(w.output_names) != tuple(HEAD_NAMES) + ("curr_bevfeat",): raise ValueError(f"outputs {w.output_names} differ from {HEAD_NAMES + ('curr_bevfeat',)}") bn = self.depth_bn() if bn.channels != cfg.cam_params_size: raise ValueError("depth_net.bn must have 27 channels") def find_onnx(weights_dir: Any = None) -> Optional[Path]: """``/bevdet_one_lt_d.onnx`` if it exists (``weights_dir`` may also be the file itself).""" if weights_dir is None: return None p = Path(weights_dir) if p.is_file(): return p q = p / ONNX_FILE return q if q.is_file() else None