# SPDX-License-Identifier: Apache-2.0 """The three deployed ONNX graphs of ``AutowareFoundation/ptv3@v4.0`` read as DATA, the way the TT port reads them: the vendored ``ttaw.weights.OnnxWeights`` (initializers addressed by the node that consumes them; nothing is executed), the module structure checked against the graph, BatchNorms folded in float64 with one final rounding **exactly where the port folds them**: - folded: embedding stem Linear + BN; seg-head ``up.proj`` / ``up.proj_skip`` Linear + BN; det-head fusion Linear + BN, ``point_proj`` MatMul + BN (eps 1e-5), every Conv2d + BN (``bev_refine`` eps 1e-5, the rest 1e-3), the prediction heads' Conv1d + BN; - **not** folded: the pooling BN (``down.norm``), which follows the segment max (a negative scale does not commute with max), so it stays an affine ``x * scale + shift`` (float64 -> float32); - kept literal: sub-manifold conv (KRSC ``[Cout, 3, 3, 3, Cin]`` -> ``[Cout, 27, Cin]``, bias), the CPE Linear and LayerNorm, qkv / proj / fc1 / fc2, LayerNorms, the det decoder (q / k / v / out projections, position MLPs, FFN). Every wiring assumption is asserted from the graph (``_check_*``): node op types, BN / LN epsilons, the ``GetIndicePairsImplicitGemm`` spatial shapes, the per-block serialization order (``Gather(.Concat, b % 2)``), the patch size 512, the head count of every attention reshape, the q scale 0.25, the CPE input of the decoder blocks (``proj_skip``, the quirk of SPEC 4.2), the det-head conv dilations / pads, the peak MaxPool on classes 0-3, the ``TopK`` K (500), the decoder ``/ 4`` scale and the ``bev_pos`` constants (``k + 0.5``, k = 0..255). numpy only (``onnx`` is imported lazily by ``OnnxWeights``); no side effects on import. """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np from ..ttaw.weights import BatchNorm, OnnxWeights, fold_bn_conv, fold_bn_linear from .config import (BN_EPS, DEC_DEPTHS, ENC_CHANNELS, ENC_DEPTHS, HEAD_DIM, LN_EPS, NUM_PROPOSALS, ONNX_FILES, PATCH_SIZE, SPATIAL_SHAPES) __all__ = ["Linear", "Affine", "Norm", "SubmConv", "Block", "EncStage", "Encoder", "DecStage", "SegHead", "Conv2d", "MHA", "DetHead", "PTv3Weights", "load_weights"] @dataclass class Linear: """``y = x @ w.T + b`` (torch layout ``w (out, in)``), float32.""" w: np.ndarray b: Optional[np.ndarray] @property def shape(self) -> Tuple[int, int]: return tuple(self.w.shape) # type: ignore[return-value] @dataclass class Affine: """An unfoldable inference BN: ``y = x * scale + shift`` (float32 rounded from float64).""" scale: np.ndarray shift: np.ndarray @dataclass class Norm: """LayerNorm over the last dim.""" g: np.ndarray b: np.ndarray eps: float @dataclass class SubmConv: """Sub-manifold 3x3x3 conv in gather form: ``out = sum_k x[nmap[:, k]] @ w[:, k, :].T + b``; ``w (Cout, 27, Cin)``.""" w: np.ndarray b: np.ndarray stage: int # which neighbour map (encoder stage) it reads @dataclass class Block: """PTv3 block: CPE (subm conv -> Linear -> LN, residual), serialized attention (LN, qkv, patches, proj, residual), MLP (LN, fc1, GELU, fc2, residual). ``order`` = serialization order index, ``heads`` = C / 16.""" name: str cpe_conv: SubmConv cpe_linear: Linear cpe_norm: Norm norm1: Norm qkv: Linear proj: Linear norm2: Norm fc1: Linear fc2: Linear order: int heads: int stage: int @dataclass class EncStage: down: Optional[Linear] # pooling projection (stages 1..4) down_norm: Optional[Affine] # BN after the segment max (not foldable) blocks: List[Block] @dataclass class Encoder: embed: Linear # Linear 4 -> 32 with the stem BN folded (then GELU) stages: List[EncStage] @dataclass class DecStage: up: Linear # Linear + BN folded (then GELU), applied to the coarser level, gathered by cluster skip: Linear # proj_skip: Linear + BN folded (then GELU) block: Optional[Block] # dec_depths [0, 0, 1, 1]; its CPE reads ``skip`` (graph quirk) @dataclass class SegHead: dec: Dict[int, DecStage] # stage -> decoder (run 3, 2, 1, 0) classifier: Linear @dataclass class Conv2d: """3x3 conv (BN folded) with dilation / padding; ``relu``; ``residual`` = add the named input before the ReLU.""" name: str w: np.ndarray b: np.ndarray dilation: int padding: int relu: bool = True @dataclass class MHA: q: Linear k: Linear v: Linear out: Linear heads: int scale_div: float # the graph divides q by this (4.0 = sqrt(16)) @dataclass class DetHead: fusion_skip: Linear # point_feat_3: Linear + BN folded, GELU fusion_proj: Linear # point_feat_4: Linear + BN folded, GELU, gathered by cluster_3 point_proj: Linear # MatMul (no bias) + BN (eps 1e-5) folded, ReLU bev_refine: Conv2d stem: Conv2d blocks: List[Tuple[Conv2d, Conv2d]] # residual blocks: relu(conv2(relu(conv1(x))) + x) shared: Conv2d heat_hidden: Conv2d heat_out: Conv2d # 128 -> 7 logits (bias, no BN, no ReLU) class_encoding: Linear # Conv1d 7 -> 128 of the one-hot (w (128, 7)) bev_pos: np.ndarray # (65536, 2) float32: (x + 0.5, y + 0.5) at k = y * 256 + x pooled_classes: Tuple[int, ...] # classes whose interior goes through the 3x3 max pool (0..3) num_proposals: int query_pos: Tuple[Linear, Linear] # position MLP 2 -> 128 -> 128 (ReLU between) key_pos: Tuple[Linear, Linear] self_attn: MHA cross_attn: MHA norm1: Norm norm2: Norm norm3: Norm ffn: Tuple[Linear, Linear] # 128 -> 256 ReLU -> 128 heads: Dict[str, Tuple[Linear, Linear]] # name -> (Conv1d + BN folded, ReLU; Conv1d + bias) PRED_HEADS = (("center", 2), ("height", 1), ("dim", 3), ("rot", 2), ("vel", 2), ("heatmap", 7)) def _f32(a: np.ndarray) -> np.ndarray: return np.ascontiguousarray(np.asarray(a, np.float32)) class _Reader: """Graph access helpers on one ``OnnxWeights``.""" def __init__(self, w: OnnxWeights): self.w = w def gemm(self, node: str) -> Linear: g = self.w.gemm(node) assert g.trans_b and not g.trans_a and g.alpha == 1.0 and g.beta == 1.0, node return Linear(_f32(g.weight), None if g.bias is None else _f32(g.bias)) def bn_after(self, node: str, eps: float) -> BatchNorm: out = self.w.node(node).outputs[0] bn = self.w.batchnorm(self.w.consumer_of(out, "BatchNormalization").name) assert abs(bn.eps - eps) < 1e-9 * max(1.0, eps) + 1e-12, (node, bn.eps, eps) return bn def gemm_bn(self, node: str, eps: float = BN_EPS) -> Linear: g = self.w.gemm(node) assert g.trans_b, node wf, bf = fold_bn_linear(g.weight, g.bias, self.bn_after(node, eps), layout="out_in") return Linear(_f32(wf), _f32(bf)) def ln(self, node: str, eps: float = LN_EPS) -> Norm: n = self.w.node(node) assert n.op_type == "LayerNormalization", node e = float(n.attrs.get("epsilon", 1e-5)) assert abs(e - eps) < 1e-9, (node, e) return Norm(_f32(self.w.param(node, 1)), _f32(self.w.param(node, 2)), e) def transposed_matmul(self, node: str) -> Tuple[np.ndarray, Optional[np.ndarray]]: """``MatMul(x, Transpose(W))`` (+ a following ``Add(bias)``) -> torch ``(W (out, in), bias)``.""" n = self.w.node(node) assert n.op_type == "MatMul", node tr = self.w.producer(n.inputs[1]) assert tr is not None and tr.op_type == "Transpose" and list(tr.attrs.get("perm", [])) == [1, 0], node weight = self.w.array(tr.inputs[0]) bias = None adds = [c for c in self.w.consumers(n.outputs[0]) if c.op_type == "Add"] if len(adds) == 1: consts = [t for t in adds[0].inputs if self.w.has(t)] if len(consts) == 1: bias = self.w.array(consts[0]) return _f32(weight), None if bias is None else _f32(bias) def mm_linear(self, node: str) -> Linear: w, b = self.transposed_matmul(node) assert b is not None, node return Linear(w, b) def _subm(r: _Reader, node: str, bias_add: str, stage: int, spatial: Tuple[int, int, int]) -> SubmConv: w = r.w n = w.node(node) assert n.op_type == "ImplicitGemm" and int(n.attrs.get("is_subm", 0)) == 1, node pairs = w.producer(n.inputs[2]) assert pairs is not None and pairs.op_type == "GetIndicePairsImplicitGemm", node assert list(pairs.attrs["spatial_shape"]) == list(spatial) and int(pairs.attrs["subm"]) == 1, (node, pairs.attrs) assert list(pairs.attrs["ksize"]) == [3, 3, 3] and list(pairs.attrs["dilation"]) == [1, 1, 1], node filt = w.param(node, 1) # KRSC [Cout, 3, 3, 3, Cin] cout, cin = filt.shape[0], filt.shape[-1] add = w.node(bias_add) assert add.op_type == "Add" and n.outputs[0] in add.inputs, bias_add b = [w.array(t) for t in add.inputs if w.has(t)] assert len(b) == 1, bias_add return SubmConv(_f32(filt.reshape(cout, 27, cin)), _f32(b[0]), stage) def _attn_checks(r: _Reader, p: str, channels: int, order: Optional[int]) -> int: """Asserts of one serialized attention ``p`` (``/enc/enc0/block0/attn``); returns the head count.""" w = r.w shape = w.param(f"{p}/Reshape", 1).tolist() assert shape[1] == PATCH_SIZE and shape[2] == 3 and shape[4] == HEAD_DIM, (p, shape) heads = int(shape[3]) assert heads * HEAD_DIM == channels, (p, heads, channels) scale = float(w.param(f"{p}/Mul_1", 1)) assert abs(scale - HEAD_DIM ** -0.5) < 1e-7, (p, scale) if order is not None: # encoder: Gather(.Concat, order) assert int(w.param(f"{p}/Gather_1", 1)) == order, p return heads def _block(r: _Reader, p: str, name: str, stage: int, order: int, spatial, *, check_order: bool = True) -> Block: channels = ENC_CHANNELS[stage] heads = _attn_checks(r, f"{p}/attn", channels, order if check_order else None) return Block( name=name, cpe_conv=_subm(r, f"{p}/cpe/cpe.0/ImplicitGemm", f"{p}/cpe/cpe.0/Add", stage, spatial), cpe_linear=r.gemm(f"{p}/cpe/cpe.1/Gemm"), cpe_norm=r.ln(f"{p}/cpe/cpe.2/LayerNormalization"), norm1=r.ln(f"{p}/norm1/norm1.0/LayerNormalization"), qkv=r.gemm(f"{p}/attn/qkv/Gemm"), proj=r.gemm(f"{p}/attn/proj/Gemm"), norm2=r.ln(f"{p}/norm2/norm2.0/LayerNormalization"), fc1=r.gemm(f"{p}/mlp/mlp.0/fc1/Gemm"), fc2=r.gemm(f"{p}/mlp/mlp.0/fc2/Gemm"), order=order, heads=heads, stage=stage) def _encoder(w: OnnxWeights) -> Encoder: r = _Reader(w) embed = r.gemm_bn("/embedding/stem/stem.0/Gemm") stages: List[EncStage] = [] for s in range(5): down = norm = None if s > 0: down = r.gemm(f"/enc/enc{s}/down/proj/Gemm") seg = w.node(f"/enc/enc{s}/down/SegmentCSR") assert seg.attrs.get("reduce") == "max", seg.attrs assert seg.inputs[1] == f"serialized_pooling_{s - 1}_indptr", seg.inputs bn = w.batchnorm(f"/enc/enc{s}/down/norm/norm.0/BatchNormalization") assert abs(bn.eps - BN_EPS) < 1e-9, bn.eps norm = Affine(_f32(bn.scale), _f32(bn.shift)) blocks = [_block(r, f"/enc/enc{s}/block{b}", f"enc{s}.block{b}", s, b % 2, SPATIAL_SHAPES[s]) for b in range(ENC_DEPTHS[s])] stages.append(EncStage(down, norm, blocks)) assert embed.w.shape == (ENC_CHANNELS[0], 4) return Encoder(embed, stages) def _seg_head(w: OnnxWeights) -> SegHead: r = _Reader(w) dec: Dict[int, DecStage] = {} for s in (3, 2, 1, 0): p = f"/seg3d_head/dec/dec{s}" up = r.gemm_bn(f"{p}/up/proj/proj.0/Gemm") skip = r.gemm_bn(f"{p}/up/proj_skip/proj_skip.0/Gemm") block = None if DEC_DEPTHS[s]: block = _block(r, f"{p}/block0", f"dec{s}.block0", s, 0, SPATIAL_SHAPES[s], check_order=False) # the quirk: the CPE sparse conv reads the proj_skip output, not the unpooled sum (SPEC 4.2) cpe_in = w.node(f"{p}/block0/cpe/cpe.0/ImplicitGemm").inputs[0] assert cpe_in == f"{p}/up/proj_skip/proj_skip.2/Mul_1_output_0", cpe_in # order 0 ("z") of the level's serialization: Gather(serialized_pooling_{s-1}_serialized_order, 0) g = w.node(f"{p}/block0/attn/Gather_1") assert g.inputs[0] == f"serialized_pooling_{s - 1}_serialized_order" and int(w.param(g.name, 1)) == 0, g dec[s] = DecStage(up, skip, block) cls = r.gemm("/seg3d_head/classifier/Gemm") assert w.find_node("/Softmax").attrs.get("axis", -1) in (1, -1) return SegHead(dec, cls) def _conv(w: OnnxWeights, r: _Reader, node: str, name: str, eps: float = BN_EPS, *, bn: bool = True, relu: bool = True) -> Conv2d: c = w.conv(node) assert c.kernel_shape == (3, 3) and c.strides == (1, 1) and c.group == 1, (node, c) d = c.dilations[0] assert c.dilations == (d, d) and c.pads == (d, d, d, d), (node, c.dilations, c.pads) if bn: wf, bf = fold_bn_conv(c.weight, c.bias, r.bn_after(node, eps)) else: wf, bf = c.weight, c.bias return Conv2d(name, _f32(wf), _f32(bf), int(d), int(d), relu) def _mha(r: _Reader, p: str) -> MHA: w = r.w div = float(w.param(f"{p}/Div", 1)) q = r.mm_linear(f"{p}/q_proj/MatMul") heads = 8 # 8 x 16 (SPEC 4.3 item 7); the ORT tests check the outputs assert abs(div - (q.w.shape[0] // heads) ** 0.5) < 1e-6, (p, div) return MHA(q, r.mm_linear(f"{p}/k_proj/MatMul"), r.mm_linear(f"{p}/v_proj/MatMul"), r.mm_linear(f"{p}/out_proj/MatMul"), heads, div) def _det_head(w: OnnxWeights) -> DetHead: r = _Reader(w) fusion_skip = r.gemm_bn("/bev_neck/fusion/proj_skip/proj_skip.0/Gemm") fusion_proj = r.gemm_bn("/bev_neck/fusion/proj/proj.0/Gemm") pw, pb = r.transposed_matmul("/bev_neck/bev_projector/point_proj/point_proj.0/MatMul") assert pb is None wf, bf = fold_bn_linear(pw, None, r.bn_after("/bev_neck/bev_projector/point_proj/point_proj.0/MatMul", 1e-5), layout="out_in") point_proj = Linear(_f32(wf), _f32(bf)) sc = w.node("/bev_neck/bev_projector/ScatterElements") assert sc.attrs.get("reduction") == "max", sc.attrs bev_refine = _conv(w, r, "/bev_neck/bev_projector/bev_refine/bev_refine.0/Conv", "bev_refine", 1e-5) stem = _conv(w, r, "/bev_neck/bev_encoder/stem/stem.0/Conv", "stem") blocks = [] for i in range(4): p = f"/bev_neck/bev_encoder/blocks/blocks.{i}" blocks.append((_conv(w, r, f"{p}/conv1/Conv", f"block{i}.conv1"), _conv(w, r, f"{p}/conv2/Conv", f"block{i}.conv2"))) assert [b[0].dilation for b in blocks] == [1, 2, 5, 2], [b[0].dilation for b in blocks] shared = _conv(w, r, "/bbox_head/shared_conv/Conv", "shared") heat_hidden = _conv(w, r, "/bbox_head/heatmap_head/heatmap_head.0/conv/Conv", "heat_hidden") heat_out = _conv(w, r, "/bbox_head/heatmap_head/heatmap_head.1/Conv", "heat_out", bn=False, relu=False) mp = w.node("/bbox_head/MaxPool") assert list(mp.attrs["kernel_shape"]) == [3, 3] and list(mp.attrs.get("pads", [0] * 4)) == [0, 0, 0, 0] pooled = tuple(int(v) for v in w.param("/bbox_head/Gather", 1).reshape(-1)) assert pooled == (0, 1, 2, 3), pooled k = int(np.asarray(w.array("/bbox_head/TopK_K")).reshape(-1)[0]) assert k == NUM_PROPOSALS, k ce = w.conv("/bbox_head/class_encoding/Conv") class_encoding = Linear(_f32(ce.weight[:, :, 0]), _f32(ce.bias)) # bev_pos: two 256-entry constants (k + 0.5 added in the graph); channel 0 varies along x (the column) xs = w.param("/bbox_head/Reshape_1", 0).reshape(-1) ys = w.param("/bbox_head/Reshape_2", 0).reshape(-1) assert np.array_equal(xs, np.arange(256, dtype=xs.dtype)) and np.array_equal(ys, np.arange(256, dtype=ys.dtype)) yy, xx = np.meshgrid(np.arange(256, dtype=np.float32), np.arange(256, dtype=np.float32), indexing="ij") bev_pos = np.stack([xx.reshape(-1) + 0.5, yy.reshape(-1) + 0.5], axis=1).astype(np.float32) dec = "/bbox_head/decoder.0" qp = (r.mm_linear(f"{dec}/query_pos_encoding/proj/proj.0/MatMul"), r.mm_linear(f"{dec}/query_pos_encoding/proj/proj.2/MatMul")) kp = (r.mm_linear(f"{dec}/key_pos_encoding/proj/proj.0/MatMul"), r.mm_linear(f"{dec}/key_pos_encoding/proj/proj.2/MatMul")) heads = {} for name, c in PRED_HEADS: p = f"/bbox_head/prediction_heads.0/{name}/{name}" h0 = w.conv(f"{p}.0/Conv") wf, bf = fold_bn_linear(h0.weight[:, :, 0], h0.bias, r.bn_after(f"{p}.0/Conv", BN_EPS), layout="out_in") h3 = w.conv(f"{p}.3/Conv") assert h3.weight.shape[0] == c, (name, h3.weight.shape) heads[name] = (Linear(_f32(wf), _f32(bf)), Linear(_f32(h3.weight[:, :, 0]), _f32(h3.bias))) return DetHead(fusion_skip, fusion_proj, point_proj, bev_refine, stem, blocks, shared, heat_hidden, heat_out, class_encoding, bev_pos, pooled, k, qp, kp, _mha(r, f"{dec}/self_attn"), _mha(r, f"{dec}/cross_attn"), r.ln(f"{dec}/norm1/LayerNormalization"), r.ln(f"{dec}/norm2/LayerNormalization"), r.ln(f"{dec}/norm3/LayerNormalization"), (r.mm_linear(f"{dec}/ffn/ffn.0/MatMul"), r.mm_linear(f"{dec}/ffn/ffn.3/MatMul")), heads) @dataclass class PTv3Weights: encoder: Encoder seg: Optional[SegHead] det: Optional[DetHead] sha256: Dict[str, str] paths: Dict[str, str] def num_parameters(self) -> int: def count(obj) -> int: if isinstance(obj, np.ndarray): return int(obj.size) if isinstance(obj, (list, tuple)): return sum(count(o) for o in obj) if isinstance(obj, dict): return sum(count(o) for o in obj.values()) if hasattr(obj, "__dataclass_fields__"): return sum(count(getattr(obj, f)) for f in obj.__dataclass_fields__ if f not in ("bev_pos",)) return 0 return count(self.encoder) + count(self.seg) + count(self.det) def load_weights(weights_dir, *, seg: bool = True, det: bool = True) -> PTv3Weights: """Read ``ptv3_encoder.onnx`` (+ seg / det heads) from ``weights_dir``.""" d = Path(weights_dir) paths = {k: str(d / v) for k, v in ONNX_FILES.items() if k == "encoder" or (k == "seg" and seg) or (k == "det" and det)} ws = {k: OnnxWeights(p) for k, p in paths.items()} return PTv3Weights(_encoder(ws["encoder"]), _seg_head(ws["seg"]) if seg else None, _det_head(ws["det"]) if det else None, {k: v.sha256 for k, v in ws.items()}, paths)