| |
| """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(<initializer>, 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" |
| |
| 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 |
| bias: Optional[np.ndarray] |
| stride: Tuple[int, int] |
| pads: Tuple[int, int, int, int] |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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]: |
| """``<weights_dir>/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 |
|
|
