code/tt_ptv3/reference/weights.py

# 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)