| |
| """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) |
|
|
|
|
| @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 |
|
|
|
|
| @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] |
| down_norm: Optional[Affine] |
| blocks: List[Block] |
|
|
|
|
| @dataclass |
| class Encoder: |
| embed: Linear |
| stages: List[EncStage] |
|
|
|
|
| @dataclass |
| class DecStage: |
| up: Linear |
| skip: Linear |
| block: Optional[Block] |
|
|
|
|
| @dataclass |
| class SegHead: |
| dec: Dict[int, DecStage] |
| 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 |
|
|
|
|
| @dataclass |
| class DetHead: |
| fusion_skip: Linear |
| fusion_proj: Linear |
| point_proj: Linear |
| bev_refine: Conv2d |
| stem: Conv2d |
| blocks: List[Tuple[Conv2d, Conv2d]] |
| shared: Conv2d |
| heat_hidden: Conv2d |
| heat_out: Conv2d |
| class_encoding: Linear |
| bev_pos: np.ndarray |
| pooled_classes: Tuple[int, ...] |
| num_proposals: int |
| query_pos: Tuple[Linear, Linear] |
| key_pos: Tuple[Linear, Linear] |
| self_attn: MHA |
| cross_attn: MHA |
| norm1: Norm |
| norm2: Norm |
| norm3: Norm |
| ffn: Tuple[Linear, Linear] |
| heads: Dict[str, Tuple[Linear, Linear]] |
|
|
|
|
| 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) |
| 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: |
| 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) |
| |
| 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 |
| |
| 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 |
| 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)) |
| |
| 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) |
|
|
