# SPDX-License-Identifier: Apache-2.0 """The deployed ``bevfusion_lidar.onnx`` of AutowareFoundation/bevfusion@v2.0 read as DATA, the way the TT port reads it: the vendored ``ttaw.weights.OnnxWeights`` (initializers and Constant tensors addressed by the node that consumes them; nothing is executed), the module structure checked against the graph topology, every BatchNormalization folded in float64 with one final rounding (the export keeps all 44 BNs unfolded; the port folds each into its conv / linear, so this reference does too). The export (PyTorch 2.2, opset 17 + domain ``autoware``) names its nodes after the mmdet3d modules. Every assumption the reference makes about the wiring is asserted from the graph (S:bevfusion:205-264): - sparse encoder (``/pts_middle_encoder``): 21 ``autoware::ImplicitGemm`` (KRSC filters ``[Cout, k0, k1, k2, Cin]``, no bias, no fused activation) each followed by ``BatchNormalization`` (eps 1e-3, **folded here**: ``W' = W * s`` along Cout, ``b' = beta - mean * s``), the basic blocks' residual ``Add`` and ``Relu``; their ``GetIndicePairsImplicitGemm`` attributes (subm, ksize, stride, padding, spatial_shape) must be the layer plan of ``host/preprocess.py`` (4 sub-manifold levels, 3 stride-2 convs, the 1x1x3 / (1, 1, 2) conv_out); the VFE (``ReduceSum / Cast / Div``), the ``coors`` flip and the to-dense ``ScatterND`` -> ``Transpose(0, 4, 3, 1, 2)`` -> ``Reshape [1, C * Z, X, Y]`` (channel ``c * 2 + z``); - SECOND (``/pts_backbone/blocks.{0,1}``): 6 + 6 Conv 3x3 (no bias) + BN (eps 1e-3, folded) + Relu, stride 2 first in block 1; SECONDFPN (``/pts_neck``): Conv 1x1 + BN + Relu, ConvTranspose 2x2 s2 + BN + Relu (folded), concat; - head (``/bbox_head``): shared conv 3x3 512 -> 128 (+bias, no activation), heatmap head (3x3 + BN eps 1e-5 (folded) + Relu, 3x3 -> 5 + bias), Sigmoid, MaxPool 3x3 s1 p0 on classes 0-3 scattered into a zero map (class 4 passes), ``heat * (heat == local_max)``, TopK ``K`` = 500 over the flattened 5 x 32,400, ``class = idx // 32400``, the class encoding Conv1d 5 -> 128 of a one-hot, the constant ``bev_pos`` ``(i + 0.5, j + 0.5)`` at ``k = i * 180 + j``, one decoder layer (position-embedding MLPs Conv1d 2 -> 128 + BN (folded) + Relu + Conv1d 128 -> 128; 8-head self- and cross-attention from ``in_proj_weight`` slices with a ``/ 4`` pre-scale; post-LayerNorms eps 1e-5; FFN 128 -> 256 -> 128) and six prediction heads (Conv1d 128 -> 64 + BN (folded) + Relu, Conv1d 64 -> C); - outputs: ``bbox_pred = concat[center + query_pos, height, dim, rot, vel]`` (10, K), ``score = max_c(sigmoid( heatmap) * query_heat * onehot(label))`` (K,) and ``label_pred = class`` (K,) int64. :attr:`BEVFusionWeights.onnx_taps` maps every reference tap name of the dense part to the ONNX tensor that holds the same value, so the ONNX Runtime golden of the plugin-free dense sub-graph can expose matching intermediates (``reference/ort.py``). The sparse encoder has no ORT counterpart (``autoware::`` plugins): its oracle is the research reference (SPEC 6.2). """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple import numpy as np from ..ttaw.weights import OnnxWeights, file_sha256, fold_bn, fold_bn_conv, fold_bn_conv_transpose from .config import BEV_SHAPE, NUM_CLASSES, ONNX_FILE, SPARSE_LEVELS, BEVFusionConfig __all__ = ["SparseLayer", "Layer", "AttentionParams", "DecoderParams", "HeadParams", "BEVFusionWeights", "load_weights", "PREDICTION_HEADS", "ENCODER_PLAN", "PME", "BB", "NK", "HB", "DEC", "DENSE_INPUT"] PME, BB, NK, HB = "/pts_middle_encoder", "/pts_backbone", "/pts_neck", "/bbox_head" DEC = f"{HB}/decoder.0" DENSE_INPUT = f"{PME}/Reshape_3_output_0" # the to-dense BEV [1, 256, 180, 180]: input of the dense sub-graph # prediction heads in ONNX output order, channels PREDICTION_HEADS = (("center", 2), ("height", 1), ("dim", 3), ("rot", 2), ("vel", 2), ("heatmap", NUM_CLASSES)) # The sparse encoder, layer by layer: (tap name, ONNX module path, neighbour map of host/preprocess.py, Cin, Cout, # GetIndicePairs attributes (subm, ksize, stride, padding, spatial_shape), residual source tap or None). def _encoder_plan() -> List[Tuple[str, str, str, int, int, tuple, Optional[str]]]: p: List[Tuple[str, str, str, int, int, tuple, Optional[str]]] = [] shapes = [SPARSE_LEVELS[k] for k in ("l1", "l2", "l3", "l4")] p.append(("sparse.conv_input", f"{PME}/conv_input/conv_input.0", "subm1", 4, 16, (1, (3, 3, 3), (1, 1, 1), (1, 1, 1), shapes[0]), None)) chans = (16, 32, 64, 128) prev = "sparse.conv_input" for s in range(1, 5): c = chans[s - 1] for b in range(2): pre = f"{PME}/encoder_layer{s}/encoder_layer{s}.{b}" attrs = (1, (3, 3, 3), (1, 1, 1), (1, 1, 1), shapes[s - 1]) p.append((f"sparse.layer{s}.{b}.conv1", f"{pre}/conv1", f"subm{s}", c, c, attrs, None)) p.append((f"sparse.layer{s}.{b}", f"{pre}/conv2", f"subm{s}", c, c, attrs, prev)) prev = f"sparse.layer{s}.{b}" if s < 4: pad = (1, 1, 0) if s == 3 else (1, 1, 1) p.append((f"sparse.down{s}", f"{PME}/encoder_layer{s}/encoder_layer{s}.2/encoder_layer{s}.2.0", f"down{s}", c, chans[s], (0, (3, 3, 3), (2, 2, 2), pad, shapes[s - 1]), None)) prev = f"sparse.down{s}" p.append(("sparse.conv_out", f"{PME}/conv_out/conv_out.0", "conv_out", 128, 128, (0, (1, 1, 3), (1, 1, 2), (0, 0, 0), shapes[3]), None)) return p ENCODER_PLAN = _encoder_plan() @dataclass(frozen=True) class SparseLayer: """One sparse conv, BN folded: ``y = relu(gather_conv(x, weight) + bias (+ residual))``.""" name: str # tap name ("sparse.layer2.1", "sparse.down3", ...) map: str # neighbour map of the host plan ("subm2", "down3", "conv_out", ...) weight: np.ndarray # float32 KRSC [Cout, k0, k1, k2, Cin], BN folded bias: np.ndarray # float32 [Cout] (the folded BN shift) residual: Optional[str] # tap name of the identity added before the Relu (block conv2), or None onnx_node: str # the ImplicitGemm node bn_node: str @property def in_channels(self) -> int: return int(self.weight.shape[-1]) @property def out_channels(self) -> int: return int(self.weight.shape[0]) @property def kernel_volume(self) -> int: return int(np.prod(self.weight.shape[1:4])) @property def num_params(self) -> int: return int(self.weight.size + self.bias.size) @dataclass(frozen=True) class Layer: """One dense conv layer, BN folded: ``y = relu?(conv(x, weight) + bias)`` (2-D, or 1-D for the head's Conv1d).""" name: str # tap name ("backbone.block0.conv0", "neck.deblock1", "heads.center.hidden", ...) kind: str # "conv" | "conv_transpose" | "conv1d" weight: np.ndarray # float32: conv [Cout, Cin, k, k]; conv_transpose [Cin, Cout, k, k]; conv1d [Cout, Cin, 1] bias: np.ndarray # float32 [Cout] stride: int padding: int relu: bool onnx_node: str onnx_output: str # ONNX tensor equal to this layer's output (after the Relu when there is one) bn_node: Optional[str] = None @property def kernel(self) -> int: return int(self.weight.shape[-1]) @property def in_channels(self) -> int: return int(self.weight.shape[0] if self.kind == "conv_transpose" else self.weight.shape[1]) @property def out_channels(self) -> int: return int(self.weight.shape[1] if self.kind == "conv_transpose" else self.weight.shape[0]) @property def num_params(self) -> int: return int(self.weight.size + self.bias.size) @dataclass(frozen=True) class AttentionParams: """``nn.MultiheadAttention`` as exported: ``q = x_q @ wq + bq`` (``[in, out]``; the rows of ``in_proj_weight`` sliced and transposed by the graph), likewise k / v, per head ``softmax((q / div) @ k^T) @ v``, ``out = concat_heads @ wo + bo`` (``wo`` stored ``[in, out]`` here; the ONNX Gemm keeps ``[out, in]`` with transB=1).""" wq: np.ndarray bq: np.ndarray wk: np.ndarray bk: np.ndarray wv: np.ndarray bv: np.ndarray wo: np.ndarray bo: np.ndarray num_heads: int head_dim: int q_divisor: float # the ONNX pre-scale Div constant (sqrt(head_dim) = 4.0) @property def scale(self) -> float: return 1.0 / self.q_divisor @dataclass(frozen=True) class DecoderParams: """One TransFusion decoder layer (post-norm): self-attention on ``query + query_pos_embed`` (Q = K = V), cross-attention ``query = x + query_pos_embed`` over ``key = value = lidar_feat + key_pos_embed``, FFN.""" self_posembed: Tuple[Layer, Layer] # Conv1d 2 -> 128 + BN (folded) + Relu, Conv1d 128 -> 128 cross_posembed: Tuple[Layer, Layer] self_attn: AttentionParams cross_attn: AttentionParams norms: Tuple[Tuple[np.ndarray, np.ndarray, float], ...] # (gamma, beta, eps) x 3 ffn_w1: np.ndarray # (128, 256) [in, out] ffn_b1: np.ndarray ffn_w2: np.ndarray # (256, 128) ffn_b2: np.ndarray @dataclass(frozen=True) class HeadParams: shared: Layer # 3x3 512 -> 128 + bias, no activation (lidar_feat) heatmap_hidden: Layer # 3x3 128 -> 128 + BN (folded) + Relu heatmap_out: Layer # 3x3 128 -> 5 (logits) nms_kernel: int # 3 (MaxPool kernel; stride 1, pad 0, zero border on pooled classes) pooled_classes: Tuple[int, ...] # (0, 1, 2, 3): class 4 (PEDESTRIAN) is not pooled num_proposals: int # TopK K = 500 class_encoding_w: np.ndarray # (128, 5) = Conv1d weight[:, :, 0] class_encoding_b: np.ndarray # (128,) bev_pos: np.ndarray # (32400, 2) float32 constant: (i + 0.5, j + 0.5), k = i * 180 + j decoder: DecoderParams prediction_heads: Tuple[Tuple[str, Layer, Layer], ...] # (name, Conv1d 128 -> 64 + BN + Relu, Conv1d 64 -> C) @property def class_table(self) -> np.ndarray: """``(5, 128)`` float32: row c = ``class_encoding(one_hot(c))`` = ``W[:, c] + b``.""" return (self.class_encoding_w.T.astype(np.float32) + self.class_encoding_b.astype(np.float32)).astype( np.float32) @dataclass(frozen=True) class BEVFusionWeights: encoder: Tuple[SparseLayer, ...] # 21 sparse convs in forward order (ENCODER_PLAN) blocks: Tuple[Tuple[Layer, ...], ...] # SECOND: 2 blocks of 6 3x3 convs deblocks: Tuple[Layer, ...] # SECONDFPN: conv 1x1, convT 2x2 s2 head: HeadParams onnx_taps: Dict[str, str] # reference tap name -> ONNX tensor name (dense part) sha256: str path: str def dense_layers(self) -> List[Layer]: """Every 2-D conv layer (backbone, neck, head) in forward order.""" h = self.head return ([layer for block in self.blocks for layer in block] + list(self.deblocks) + [h.shared, h.heatmap_hidden, h.heatmap_out]) def num_params(self) -> int: """Learned parameters after the BN folds (the ONNX has 8,308,692 float32 weights incl. the BN statistics).""" h = self.head d = h.decoder n = sum(x.num_params for x in self.encoder) + sum(x.num_params for x in self.dense_layers()) n += h.class_encoding_w.size + h.class_encoding_b.size for pe in (d.self_posembed, d.cross_posembed): n += sum(x.num_params for x in pe) for a in (d.self_attn, d.cross_attn): n += sum(x.size for x in (a.wq, a.bq, a.wk, a.bk, a.wv, a.bv, a.wo, a.bo)) n += sum(g.size + b.size for g, b, _ in d.norms) n += sum(x.size for x in (d.ffn_w1, d.ffn_b1, d.ffn_w2, d.ffn_b2)) for _, hid, fin in h.prediction_heads: n += hid.num_params + fin.num_params return int(n) # ------------------------------------------------------------------------------------------------- helpers def _f32(a: Any) -> np.ndarray: return np.asarray(a, dtype=np.float32) def _single(found: Sequence[Any], what: str) -> Any: if len(found) != 1: raise ValueError(f"expected exactly one {what}, found {len(found)}") return found[0] def _expect(cond: bool, msg: str) -> None: if not cond: raise ValueError(f"bevfusion_lidar.onnx structure: {msg}") def _consumer(w: OnnxWeights, tensor: str, op_type: str) -> Any: cons = [n for n in w.consumers(tensor) if n.op_type == op_type] return _single(cons, f"{op_type} consumer of {tensor!r}") def _scalar(w: OnnxWeights, tensor: str) -> float: a = np.asarray(w.array(tensor)).reshape(-1) _expect(a.size == 1, f"{tensor} is not a scalar constant") return float(a[0]) def _fold_relu_chain(w: OnnxWeights, node_name: str) -> Tuple[Any, Any, Any]: """``node -> BatchNormalization -> Relu`` (the single consumers): returns (node, bn, relu).""" node = w.node(node_name) bn = _consumer(w, node.outputs[0], "BatchNormalization") relu = _consumer(w, bn.outputs[0], "Relu") return node, bn, relu def _conv_bn_relu(w: OnnxWeights, name: str, node_name: str, *, kind: str = "conv") -> Layer: """Conv / ConvTranspose (no bias) + BatchNormalization (folded) + Relu.""" node, bn, relu = _fold_relu_chain(w, node_name) c = w.conv(node_name) _expect(c.group == 1 and all(d == 1 for d in c.dilations) and len(set(c.strides)) == 1 and len(set(c.pads)) <= 1 and len(set(c.kernel_shape)) == 1 and c.auto_pad == "NOTSET", f"{node_name}: unexpected attributes {c}") b = None if c.bias is None else np.asarray(c.bias, np.float64) if kind == "conv_transpose": _expect(tuple(c.kernel_shape) == tuple(c.strides) and not any(c.pads), f"{node_name}: ConvTranspose k != s") wf, bf = fold_bn_conv_transpose(np.asarray(c.weight, np.float64), b, w.batchnorm(bn.name)) else: wf, bf = fold_bn_conv(np.asarray(c.weight, np.float64), b, w.batchnorm(bn.name)) pad = int(c.pads[0]) if c.pads else 0 return Layer(name, kind, _f32(wf), _f32(bf), int(c.strides[0]), pad, True, node_name, relu.outputs[0], bn_node=bn.name) def _conv_plain(w: OnnxWeights, name: str, node_name: str, *, kind: str = "conv") -> Layer: """Conv with bias, no BN, no activation.""" c = w.conv(node_name) _expect(c.group == 1 and all(d == 1 for d in c.dilations) and len(set(c.strides)) == 1 and len(set(c.pads)) <= 1, f"{node_name}: unexpected attributes {c}") bias = np.zeros(c.weight.shape[0], np.float32) if c.bias is None else _f32(c.bias) pad = int(c.pads[0]) if c.pads else 0 return Layer(name, kind, _f32(c.weight), bias, int(c.strides[0]), pad, False, node_name, w.node(node_name).outputs[0]) # ----------------------------------------------------------------------------------------- sparse encoder def _load_encoder(w: OnnxWeights) -> Tuple[Tuple[SparseLayer, ...], Dict[str, str]]: # VFE and coordinate flip feeding the first conv: features = ReduceSum(voxels, 1) / float(num_points) rs = w.node("/ReduceSum") div = w.node("/Div") _expect(rs.inputs[0] == "voxels" and int(rs.attrs.get("keepdims", 1)) == 0 and list(np.asarray(w.array(rs.inputs[1])).reshape(-1)) == [1] and div.inputs[0] == rs.outputs[0], "VFE is not ReduceSum(voxels, axis=1) / num_points") sl = w.node("/Slice") _expect(sl.inputs[0] == "coors" and int(np.asarray(w.array(sl.inputs[3])).reshape(-1)[0]) == -1, "coors are not flipped (z, y, x) -> (x, y, z)") layers: List[SparseLayer] = [] taps: Dict[str, str] = {} expected_input = div.outputs[0] relu_out: Dict[str, str] = {} for name, path, map_name, cin, cout, attrs, residual in ENCODER_PLAN: gemm = _single(w.nodes("ImplicitGemm", f"{path}/ImplicitGemm"), f"ImplicitGemm of {path}") pairs = w.node(f"{path}/GetIndicePairsImplicitGemm") a = pairs.attrs got = (int(a["subm"]), tuple(int(v) for v in a["ksize"]), tuple(int(v) for v in a["stride"]), tuple(int(v) for v in a["padding"]), tuple(int(v) for v in a["spatial_shape"])) _expect(got == attrs, f"{pairs.name}: attributes {got} != the plan {attrs}") _expect(tuple(int(v) for v in a["dilation"]) == (1, 1, 1) and int(a.get("transpose", 0)) == 0 and int(a.get("batch_size", 1)) == 1, f"{pairs.name}: dilation / transpose / batch") _expect(int(gemm.attrs.get("is_subm", -1)) == attrs[0] and float(gemm.attrs.get("output_scale", 1.0)) == 1.0 and float(gemm.attrs.get("output_add_scale", 0.0)) == 0.0 and len(gemm.inputs) == 5, f"{gemm.name}: fused bias / activation / scale") _expect(gemm.inputs[0] == expected_input, f"{gemm.name} does not read the previous layer ({expected_input})") wt = np.asarray(w.array(gemm.inputs[1]), np.float64) _expect(wt.shape == (cout, *attrs[1], cin), f"{gemm.name}: filter {wt.shape} != KRSC {(cout, *attrs[1], cin)}") bn = _consumer(w, gemm.outputs[0], "BatchNormalization") bnp = w.batchnorm(bn.name) _expect(abs(bnp.eps - 1e-3) < 1e-9, f"{bn.name}: eps {bnp.eps}") wf, bf = fold_bn(wt, None, bnp, axis=0) nxt = bn.outputs[0] if residual is not None: add = _consumer(w, nxt, "Add") _expect(set(add.inputs) == {nxt, relu_out[residual]}, f"{add.name}: residual is not {residual}") nxt = add.outputs[0] relu = _consumer(w, nxt, "Relu") layers.append(SparseLayer(name, map_name, _f32(wf), _f32(bf), residual, gemm.name, bn.name)) relu_out[name] = relu.outputs[0] taps[name] = relu.outputs[0] expected_input = relu.outputs[0] # to-dense: ScatterND of the conv_out features into [1, X, Y, Z, C], Transpose(0, 4, 3, 1, 2), Reshape scat = _single(w.nodes("ScatterND", f"{PME}/*"), "to-dense ScatterND") _expect(w.producer(scat.inputs[2]) is not None, "to-dense updates") t1 = w.node(f"{PME}/Transpose_1") _expect(t1.inputs[0] == scat.outputs[0] and list(t1.attrs["perm"]) == [0, 4, 3, 1, 2], "to-dense transpose is not (0, 4, 3, 1, 2): channel order c * Z + z") rs3 = w.node(f"{PME}/Reshape_3") _expect(rs3.inputs[0] == t1.outputs[0] and rs3.outputs[0] == DENSE_INPUT, "to-dense reshape") _expect(len(layers) == 21, f"{len(layers)} sparse convs, expected 21") taps["sparse.bev"] = DENSE_INPUT return tuple(layers), taps # ------------------------------------------------------------------------------------- SECOND / SECONDFPN def _load_backbone_neck(w: OnnxWeights) -> Tuple[tuple, tuple, Dict[str, str]]: taps: Dict[str, str] = {} blocks, prev = [], DENSE_INPUT for b in range(2): convs = w.nodes("Conv", f"{BB}/blocks.{b}/*") layers = [] for j, conv in enumerate(convs): _expect(conv.inputs[0] == prev, f"{conv.name} does not consume the previous layer output") layer = _conv_bn_relu(w, f"backbone.block{b}.conv{j}", conv.name) stride = 2 if (b == 1 and j == 0) else 1 _expect(layer.kernel == 3 and layer.padding == 1 and layer.stride == stride and not np.asarray(w.conv(conv.name).bias if w.conv(conv.name).bias is not None else 0).any(), f"{conv.name}: SECOND conv must be 3x3 pad 1 (stride {stride}), no bias") layers.append(layer) taps[layer.name] = layer.onnx_output prev = layer.onnx_output blocks.append(tuple(layers)) taps[f"backbone.block{b}"] = prev _expect([len(b) for b in blocks] == [6, 6], f"blocks have {[len(b) for b in blocks]} convs, expected 6 / 6") concat = w.node(f"{NK}/Concat") _expect(int(concat.attrs.get("axis", -1)) == 1 and len(concat.inputs) == 2, "neck concat") d0 = _conv_bn_relu(w, "neck.deblock0", f"{NK}/deblocks.0/deblocks.0.0/Conv") d1 = _conv_bn_relu(w, "neck.deblock1", f"{NK}/deblocks.1/deblocks.1.0/ConvTranspose", kind="conv_transpose") _expect(list(concat.inputs) == [d0.onnx_output, d1.onnx_output], "neck concat inputs") _expect(w.node(d0.onnx_node).inputs[0] == taps["backbone.block0"] and w.node(d1.onnx_node).inputs[0] == taps["backbone.block1"], "deblock inputs") _expect((d0.kind, d0.kernel, d0.stride, d1.kind, d1.kernel, d1.stride) == ("conv", 1, 1, "conv_transpose", 2, 2), "deblocks (expected conv 1x1, convT 2x2 s2)") taps.update({"neck.deblock0": d0.onnx_output, "neck.deblock1": d1.onnx_output, "neck": concat.outputs[0]}) return tuple(blocks), (d0, d1), taps # ---------------------------------------------------------------------------------------------------- head def _attention(w: OnnxWeights, prefix: str) -> AttentionParams: a = f"{prefix}/attn" mats = [] for i, suffix in enumerate(("", "_1", "_2")): mm = w.node(f"{a}/MatMul{suffix}") tr = w.producer(mm.inputs[1]) sl = w.producer(tr.inputs[0]) _expect(tr.op_type == "Transpose" and sl.op_type == "Slice", f"{mm.name}: weight is not Transpose(Slice)") mats.append(sl.inputs[0]) _expect(len(set(mats)) == 1, f"{a}: Q / K / V do not slice one in_proj_weight") W = np.asarray(w.array(mats[0]), np.float32) # [3 * E, E] E = W.shape[1] _expect(W.shape == (3 * E, E), f"{a}: in_proj_weight {W.shape}") bias_name = w.producer(w.node(f"{a}/Add_2").inputs[0]).inputs[0] B = np.asarray(w.array(bias_name), np.float32) _expect(B.shape == (3 * E,), f"{a}: in_proj_bias {B.shape}") for n, mm in (("Add_2", "MatMul"), ("Add_3", "MatMul_1"), ("Add_4", "MatMul_2")): _expect(w.node(f"{a}/{mm}").outputs[0] in w.node(f"{a}/{n}").inputs, f"{a}/{n} does not add the bias of {mm}") g = w.gemm(f"{a}/Gemm") _expect(g.trans_b and not g.trans_a and g.alpha == 1.0 and g.beta == 1.0, f"{a}/Gemm attributes") heads = int(round(_scalar(w, w.node(f"{a}/Div").inputs[1]))) hd = E // heads div = _scalar(w, w.node(f"{a}/Div_3").inputs[1]) _expect(heads * hd == E and abs(div - np.sqrt(hd)) < 1e-6, f"{a}: {heads} heads, divisor {div}") return AttentionParams(W[:E].T.copy(), B[:E].copy(), W[E:2 * E].T.copy(), B[E:2 * E].copy(), W[2 * E:].T.copy(), B[2 * E:].copy(), _f32(g.weight).T.copy(), _f32(g.bias), heads, hd, div) def _posembed(w: OnnxWeights, which: str) -> Tuple[Layer, Layer]: p = f"{DEC}/{which}/position_embedding_head" node, bn, relu = _fold_relu_chain(w, f"{p}/position_embedding_head.0/Conv") c = w.conv(node.name) wf, bf = fold_bn_conv(np.asarray(c.weight, np.float64), np.asarray(c.bias, np.float64), w.batchnorm(bn.name)) l0 = Layer(f"decoder.{which}.0", "conv1d", _f32(wf), _f32(bf), 1, 0, True, node.name, relu.outputs[0], bn_node=bn.name) l1 = _conv_plain(w, f"decoder.{which}.1", f"{p}/position_embedding_head.3/Conv", kind="conv1d") _expect(w.node(l1.onnx_node).inputs[0] == l0.onnx_output and l0.weight.shape == (128, 2, 1) and l1.weight.shape == (128, 128, 1), f"{which} MLP shape / wiring") return l0, l1 def _load_head(w: OnnxWeights, cfg: Optional[BEVFusionConfig], neck_out: str) -> Tuple[HeadParams, Dict[str, str]]: taps: Dict[str, str] = {} shared = _conv_plain(w, "head.shared", f"{HB}/shared_conv/Conv") _expect(w.node(shared.onnx_node).inputs[0] == neck_out and (shared.in_channels, shared.out_channels, shared.kernel) == (512, 128, 3), "shared conv") cast1 = w.node(f"{HB}/Cast_1") _expect(cast1.inputs[0] == shared.onnx_output, "heatmap head input is not the shared conv") hh = _conv_bn_relu(w, "head.heatmap.hidden", f"{HB}/heatmap_head/heatmap_head.0/conv/Conv") _expect(abs(w.batchnorm(hh.bn_node).eps - 1e-5) < 1e-9, "heatmap head BN eps != 1e-5") ho = _conv_plain(w, "head.heatmap", f"{HB}/heatmap_head/heatmap_head.1/Conv") _expect(w.node(hh.onnx_node).inputs[0] == cast1.outputs[0] and w.node(ho.onnx_node).inputs[0] == hh.onnx_output and ho.out_channels == NUM_CLASSES and hh.kernel == ho.kernel == 3, "heatmap head") sig = w.node(f"{HB}/Sigmoid") _expect(sig.inputs[0] == ho.onnx_output, "sigmoid of the heatmap") # local max: MaxPool 3x3 s1 p0 of the pooled classes, scattered into a zero map; the other classes pass g2 = w.node(f"{HB}/Gather_2") pooled = tuple(int(v) for v in np.asarray(w.array(g2.inputs[1])).reshape(-1)) pool = w.node(f"{HB}/MaxPool") k = list(pool.attrs.get("kernel_shape", [])) _expect(g2.inputs[0] == sig.outputs[0] and pool.inputs[0] == g2.outputs[0] and len(set(k)) == 1 and list(pool.attrs.get("strides", [1, 1])) == [1, 1] and not any(pool.attrs.get("pads", [0, 0, 0, 0])), f"MaxPool wiring / attributes {pool.attrs}") g6 = w.node(f"{HB}/Gather_6") passing = tuple(int(v) for v in np.asarray(w.array(g6.inputs[1])).reshape(-1)) _expect(pooled == (0, 1, 2, 3) and passing == (4,) and g6.inputs[0] == sig.outputs[0], f"pooled classes {pooled}, passing {passing}") eq = w.node(f"{HB}/Equal_6") mul6 = w.node(f"{HB}/Mul_6") scat1 = w.node(f"{HB}/ScatterND_1") _expect(list(eq.inputs) == [sig.outputs[0], scat1.outputs[0]] and sig.outputs[0] in mul6.inputs, "heat * (heat == local_max) wiring") topk = w.node(f"{HB}/TopK") K = int(np.asarray(w.array(topk.inputs[1])).reshape(-1)[0]) _expect(int(topk.attrs.get("largest", 1)) == 1 and int(topk.attrs.get("sorted", 1)) == 1, "TopK attributes") enc = w.conv(f"{HB}/class_encoding/Conv") _expect(enc.weight.shape == (128, NUM_CLASSES, 1), f"class encoding {enc.weight.shape}") onehot = w.node(f"{HB}/OneHot") _expect(int(np.asarray(w.array(onehot.inputs[1])).reshape(-1)[0]) == NUM_CLASSES, "OneHot depth") bev_pos = _f32(w.array(f"{HB}/Constant_2_output_0")).reshape(-1, 2) H, W_ = BEV_SHAPE kk = np.arange(H * W_) _expect(bev_pos.shape == (H * W_, 2) and np.array_equal(bev_pos[:, 0], (kk // W_).astype(np.float32) + np.float32(0.5)) and np.array_equal(bev_pos[:, 1], (kk % W_).astype(np.float32) + np.float32(0.5)), "bev_pos is not (i + 0.5, j + 0.5) with k = i * 180 + j") if cfg is not None: _expect(K == cfg.num_proposals, f"ONNX TopK K={K} != ml_package num_proposals {cfg.num_proposals}") # decoder norms = [] for i in range(3): n = f"{DEC}/norms.{i}" eps = _scalar(w, w.node(f"{n}/Add").inputs[1]) norms.append((_f32(w.param(f"{n}/Mul", 1)), _f32(w.param(f"{n}/Add_1", 1)), eps)) ffn0 = w.node(f"{DEC}/ffn/layers/layers.0/layers.0.0/MatMul") ffn1 = w.node(f"{DEC}/ffn/layers/layers.1/MatMul") dec = DecoderParams( self_posembed=_posembed(w, "self_posembed"), cross_posembed=_posembed(w, "cross_posembed"), self_attn=_attention(w, f"{DEC}/self_attn"), cross_attn=_attention(w, f"{DEC}/cross_attn"), norms=tuple(norms), ffn_w1=_f32(w.array(w.producer(ffn0.inputs[1]).inputs[0])).T.copy(), ffn_b1=_f32(w.param(f"{DEC}/ffn/layers/layers.0/layers.0.0/Add", 0)), ffn_w2=_f32(w.array(w.producer(ffn1.inputs[1]).inputs[0])).T.copy(), ffn_b2=_f32(w.param(f"{DEC}/ffn/layers/layers.1/Add", 0))) _expect(all(abs(eps - 1e-5) < 1e-9 for _, _, eps in dec.norms), "LayerNorm eps != 1e-5") _expect(dec.ffn_w1.shape == (128, 256) and dec.ffn_w2.shape == (256, 128), "FFN shapes") _expect(w.node(f"{DEC}/ffn/Add").inputs[0] == w.node(f"{DEC}/norms.1/Add_1").outputs[0], "FFN residual") # wiring of the decoder inputs: Q = K = V = query + pos in self-attention, cross key = lidar_feat + key pos q_in = {w.node(f"{DEC}/self_attn/Add").inputs, w.node(f"{DEC}/self_attn/Add_1").inputs, w.node(f"{DEC}/Add").inputs} _expect(len({frozenset(x) for x in q_in}) == 1, "self-attention Q / K / V are not all query + query_pos_embed") kv = w.node(f"{DEC}/Add_1") kpe = w.node(f"{DEC}/Transpose_1").outputs[0] _expect(kpe in kv.inputs and w.node(f"{DEC}/Transpose_3").outputs[0] in kv.inputs, "cross-attention key / value are not lidar_feat + key_pos_embed") heads = [] for name, ch in PREDICTION_HEADS: p = f"{HB}/prediction_heads.0/{name}" node, bn, relu = _fold_relu_chain(w, f"{p}/{name}.0/conv/Conv") c = w.conv(node.name) b = None if c.bias is None else np.asarray(c.bias, np.float64) wf, bf = fold_bn_conv(np.asarray(c.weight, np.float64), b, w.batchnorm(bn.name)) hid = Layer(f"heads.{name}.hidden", "conv1d", _f32(wf), _f32(bf), 1, 0, True, node.name, relu.outputs[0], bn_node=bn.name) fin = _conv_plain(w, f"heads.{name}", f"{p}/{name}.1/Conv", kind="conv1d") _expect(hid.weight.shape == (64, 128, 1) and fin.weight.shape == (ch, 64, 1) and w.node(fin.onnx_node).inputs[0] == hid.onnx_output, f"prediction head {name}") heads.append((name, hid, fin)) taps[fin.name] = fin.onnx_output head = HeadParams(shared=shared, heatmap_hidden=hh, heatmap_out=ho, nms_kernel=int(k[0]), pooled_classes=pooled, num_proposals=K, class_encoding_w=_f32(enc.weight[:, :, 0]), class_encoding_b=_f32(enc.bias), bev_pos=bev_pos, decoder=dec, prediction_heads=tuple(heads)) taps.update({ "head.shared": shared.onnx_output, "head.heatmap.hidden": hh.onnx_output, "head.heatmap": ho.onnx_output, "head.heat": sig.outputs[0], "head.heat_nms": mul6.outputs[0], "head.topk_values": topk.outputs[0], "head.topk_indices": topk.outputs[1], "query.feat": w.node(f"{HB}/Add_4").outputs[0], "query.pos": w.node(f"{HB}/GatherElements_1").outputs[0], "query.heat": w.node(f"{HB}/GatherElements_2").outputs[0], "decoder.key_pos_embed": kpe, "decoder.query_pos_embed": w.node(f"{DEC}/Transpose").outputs[0], "decoder.self_attn": w.node(f"{DEC}/norms.0/Add_1").outputs[0], "decoder.cross_attn": w.node(f"{DEC}/norms.1/Add_1").outputs[0], "decoder.out": w.node(f"{DEC}/norms.2/Add_1").outputs[0], }) _check_outputs(w, heads) return head, taps def _check_outputs(w: OnnxWeights, heads: Sequence[Tuple[str, Layer, Layer]]) -> None: """``bbox_pred = concat[center + query_pos, height, dim, rot, vel]`` (batch 0), ``score = ReduceMax_c(sigmoid( heatmap) * query_heat * onehot(label))`` (batch 0), ``label_pred = class`` (batch 0).""" out = {name: fin.onnx_output for name, _, fin in heads} def through(t: str) -> str: # follow single-input Concats and the batch-0 Gather forward while True: cons = [n for n in w.consumers(t) if (n.op_type == "Concat" and len(n.inputs) == 1) or (n.op_type == "Gather" and n.inputs[1] == "/Constant_output_0")] if not cons: return t t = cons[0].outputs[0] center_add = w.node(f"{HB}/Add_5") qpos_t = w.node(f"{HB}/Transpose_2") _expect(list(center_add.inputs) == [out["center"], qpos_t.outputs[0]] and qpos_t.inputs[0] == w.node(f"{HB}/GatherElements_1").outputs[0], "center is not offset by query_pos") bbox = w.producer("bbox_pred") _expect(bbox.op_type == "Concat" and list(bbox.inputs) == [through(center_add.outputs[0]), through(out["height"]), through(out["dim"]), through(out["rot"]), through(out["vel"])], f"bbox_pred = {bbox.inputs}") sig = w.node("/Sigmoid") mul = w.node("/Mul") _expect(w.producer(sig.inputs[0]).inputs[0] == out["heatmap"] and list(mul.inputs) == [sig.outputs[0], w.node(f"{HB}/GatherElements_2").outputs[0]], "score is not sigmoid(heatmap) * query_heat") rmax = w.producer("score") _expect(rmax.op_type == "ReduceMax" and list(rmax.attrs.get("axes", [])) == [0], "score ReduceMax over classes") lab = w.producer("label_pred") _expect(lab.op_type == "Gather" and lab.inputs[0] == w.node(f"{HB}/Cast_8").outputs[0], "label_pred = class") g2 = w.node(f"{HB}/GatherElements_2") _expect(w.producer(g2.inputs[0]).inputs[0] == w.node(f"{HB}/Mul_6").outputs[0], "query_heat is not gathered from the local-max-masked heatmap (all classes)") # ---------------------------------------------------------------------------------------------------- load def load_weights(weights_dir: Any, config: Optional[BEVFusionConfig] = None) -> BEVFusionWeights: """``/bevfusion_lidar.onnx`` (or the ``.onnx`` path itself) -> :class:`BEVFusionWeights`, checked against ``config`` (K) when given.""" p = Path(weights_dir) path = p if p.suffix == ".onnx" else p / ONNX_FILE w = OnnxWeights(path) _expect(tuple(w.input_names) == ("voxels", "coors", "num_points_per_voxel"), f"inputs {w.input_names}") _expect(tuple(w.output_names) == ("bbox_pred", "score", "label_pred"), f"outputs {w.output_names}") encoder, taps = _load_encoder(w) blocks, deblocks, t2 = _load_backbone_neck(w) taps.update(t2) head, t3 = _load_head(w, config, taps["neck"]) taps.update(t3) taps.update({"bbox_pred": "bbox_pred", "score": "score", "label_pred": "label_pred"}) return BEVFusionWeights(encoder=encoder, blocks=blocks, deblocks=deblocks, head=head, onnx_taps=taps, sha256=file_sha256(path), path=str(path))