| |
| """fp32 torch rebuild of the deployed PTv3 v4.0 graphs (encoder, seg head, det head) on the host tables of |
| ``tt_ptv3.host.preprocess`` (identical to what the device consumes), with per-module taps. |
| |
| Semantics reproduced from the graphs (asserted in ``weights.py``; research/ptv3/SPEC.md 4): |
| |
| - GELU is the exact erf form; LayerNorm eps 1e-5; BN folded where ``weights.py`` says; |
| - sub-manifold conv = ``sum_k x[nmap[:, k]] @ W_k`` over the 27 taps of the host neighbour map (missing taps = 0); |
| - serialized attention: block ``b`` uses order ``b % 2``; ``Npad = ceil(N / 512) * 512``, the sequence row ``j`` |
| is voxel ``order[min(j, N - 1)]`` (pad tokens are copies of the last voxel in curve order, no mask), heads of 16, |
| ``softmax(0.25 * q k^T) v`` per 512-row patch, rows back through ``inverse``; |
| - pooling: ``proj`` Linear -> segment max over the CSR children -> BN affine -> GELU; |
| - seg decoder: ``x = GELU(skip(pf_s)) + GELU(up(x))[cluster_s]``; at stages 3 / 2 one block whose CPE reads |
| ``GELU(skip(pf_s))`` (the export quirk), order 0; classifier -> softmax -> first-max argmax; |
| - det head: fusion, ReLU point projection, BEV scatter-max into a zero 256 x 256 canvas as a gather + max over the |
| host's per-cell candidate lists (``host.preprocess.bev_index``; exact: max of non-negative values with a zero |
| row), 12 dense 3x3 convs, heat logits, sigmoid, 3x3 local max on the interior of classes 0-3 with the heat itself |
| as the seed (border cells and classes 4-6 pass), ``heat * (heat == local_max)``, ONNX-semantics TopK 500 over the |
| class-major flattening (ties: lower index first; ``topk_indices=`` teacher-forces it), query features + class |
| encoding, ``query_pos = bev_pos[pos]``, one post-norm decoder layer (self-attention with Q = K = V = query + pos |
| MLP; cross-attention to the 65,536 BEV keys + key pos MLP), FFN, six prediction heads (``center += query_pos``). |
| |
| Taps: pass ``taps={}`` (a dict, filled with float32 numpy arrays) and optionally ``tap_filter`` (a callable on the |
| name); names are listed in :data:`TAP_NAMES` (``{s}`` / ``{b}`` / ``{i}`` expanded). |
| """ |
| from __future__ import annotations |
|
|
| from typing import Callable, Dict, List, Optional, Sequence |
|
|
| import numpy as np |
|
|
| from .config import HEAD_DIM, PATCH_SIZE |
| from .weights import Block, Conv2d, DetHead, Encoder, Linear, MHA, Norm, SegHead, SubmConv |
|
|
| __all__ = ["TAP_NAMES", "encoder_forward", "seg_head_forward", "det_head_forward", "topk_onnx", "ROW_CHUNK", |
| "PATCH_CHUNK"] |
|
|
| TAP_NAMES = ( |
| "embed", "enc{s}.down", "enc{s}.block{b}.cpe", "enc{s}.block{b}.attn", "enc{s}.block{b}", "point_feat_{s}", |
| "dec{s}.up", "dec{s}.skip", "dec{s}.sum", "dec{s}.block0", "seg.logits", "seg.probs", "seg.labels", |
| "det.fusion", "det.point_proj", "det.bev", "det.bev_refine", "det.stem", "det.block{i}", "det.shared", |
| "det.dense_heatmap", "det.heat", "det.heat_nms", "det.topk_indices", "det.query_feat", "det.query_pos", |
| "det.query_pos_embed", "det.key_pos_embed", "det.self", "det.norm1", "det.cross", "det.norm2", "det.ffn", |
| "det.norm3", "det.head.{name}") |
| ROW_CHUNK = 32768 |
| PATCH_CHUNK = 32 |
|
|
|
|
| def _torch(): |
| import torch |
| return torch |
|
|
|
|
| class _Taps: |
| def __init__(self, taps: Optional[Dict[str, np.ndarray]], tap_filter: Optional[Callable[[str], bool]]): |
| self.taps, self.filter = taps, tap_filter |
|
|
| def __call__(self, name: str, x): |
| if self.taps is not None and (self.filter is None or self.filter(name)): |
| self.taps[name] = x.detach().cpu().numpy().astype(np.float32, copy=True) if hasattr(x, "detach") \ |
| else np.asarray(x) |
| return x |
|
|
|
|
| def _t(a): |
| torch = _torch() |
| a = np.ascontiguousarray(a) |
| return torch.from_numpy(a if a.flags.writeable else a.copy()) |
|
|
|
|
| def linear(x, p: Linear): |
| torch = _torch() |
| return torch.nn.functional.linear(x, _t(p.w), None if p.b is None else _t(p.b)) |
|
|
|
|
| def layer_norm(x, n: Norm): |
| torch = _torch() |
| return torch.nn.functional.layer_norm(x, (x.shape[-1],), _t(n.g), _t(n.b), n.eps) |
|
|
|
|
| def gelu(x): |
| return _torch().nn.functional.gelu(x, approximate="none") |
|
|
|
|
| def subm_conv(x, nmap: np.ndarray, conv: SubmConv): |
| """Gather-form sub-manifold conv (zero row for missing taps), float32.""" |
| torch = _torch() |
| n, cin = x.shape |
| cout = conv.w.shape[0] |
| xz = torch.cat([x, torch.zeros(1, cin, dtype=x.dtype)], 0) |
| idx = torch.from_numpy(np.where(nmap >= 0, nmap, n).astype(np.int64)) |
| wmat = _t(conv.w).permute(1, 2, 0).reshape(27 * cin, cout) |
| out = torch.empty(n, cout, dtype=x.dtype) |
| for r0 in range(0, n, ROW_CHUNK): |
| g = xz[idx[r0:r0 + ROW_CHUNK]].reshape(-1, 27 * cin) |
| out[r0:r0 + ROW_CHUNK] = g @ wmat |
| return out + _t(conv.b) |
|
|
|
|
| def serialized_attention(x, blk: Block, order: np.ndarray, inverse: np.ndarray): |
| """LN'd input ``x`` -> attention output before ``proj``'s residual (``proj`` applied).""" |
| torch = _torch() |
| n, c = x.shape |
| h = blk.heads |
| npad = -(-n // PATCH_SIZE) * PATCH_SIZE |
| seq = torch.from_numpy(np.asarray(order, np.int64)[np.minimum(np.arange(npad), n - 1)]) |
| qkv = linear(x, blk.qkv)[seq].reshape(-1, PATCH_SIZE, 3, h, HEAD_DIM).permute(2, 0, 3, 1, 4) |
| q, k, v = qkv[0], qkv[1], qkv[2] |
| out = torch.empty(q.shape[0], PATCH_SIZE, h, HEAD_DIM, dtype=x.dtype) |
| for p0 in range(0, q.shape[0], PATCH_CHUNK): |
| sl = slice(p0, p0 + PATCH_CHUNK) |
| a = torch.softmax((q[sl] * (HEAD_DIM ** -0.5)) @ k[sl].transpose(-2, -1), dim=-1) |
| out[sl] = (a @ v[sl]).transpose(1, 2) |
| o = out.reshape(-1, c)[torch.from_numpy(np.asarray(inverse, np.int64))] |
| return linear(o, blk.proj) |
|
|
|
|
| def block_forward(x, blk: Block, nmap, order, inverse, tap: _Taps, cpe_in=None, name: str = ""): |
| cpe_in = x if cpe_in is None else cpe_in |
| x = x + layer_norm(linear(subm_conv(cpe_in, nmap, blk.cpe_conv), blk.cpe_linear), blk.cpe_norm) |
| tap(f"{name}.cpe", x) |
| x = x + serialized_attention(layer_norm(x, blk.norm1), blk, order, inverse) |
| tap(f"{name}.attn", x) |
| x = x + linear(gelu(linear(layer_norm(x, blk.norm2), blk.fc1)), blk.fc2) |
| return tap(name, x) |
|
|
|
|
| def segment_max(x, indptr: np.ndarray): |
| torch = _torch() |
| indptr = np.asarray(indptr, np.int64) |
| lens = torch.from_numpy(indptr[1:] - indptr[:-1]) |
| seg = torch.repeat_interleave(torch.arange(lens.shape[0]), lens) |
| out = torch.full((lens.shape[0], x.shape[1]), -float("inf"), dtype=x.dtype) |
| return out.scatter_reduce(0, seg[:, None].expand(-1, x.shape[1]), x[:int(indptr[-1])], reduce="amax", |
| include_self=True) |
|
|
|
|
| def encoder_forward(enc: Encoder, frame, *, taps: Optional[Dict[str, np.ndarray]] = None, |
| tap_filter: Optional[Callable[[str], bool]] = None) -> List: |
| """``frame`` (``host.preprocess.Frame``) -> ``[point_feat_0..4]`` (torch float32).""" |
| torch = _torch() |
| tap = _Taps(taps, tap_filter) |
| x = tap("embed", gelu(linear(torch.from_numpy(frame.voxels.feat), enc.embed))) |
| outs = [] |
| for s, st in enumerate(enc.stages): |
| lv = frame.levels[s] |
| if s > 0: |
| pooled = segment_max(linear(x, st.down), frame.pools[s - 1].indptr) |
| x = tap(f"enc{s}.down", gelu(pooled * _t(st.down_norm.scale) + _t(st.down_norm.shift))) |
| for b, blk in enumerate(st.blocks): |
| x = block_forward(x, blk, lv.nmap, lv.order[blk.order], lv.inverse[blk.order], tap, name=f"enc{s}.block{b}") |
| outs.append(tap(f"point_feat_{s}", x)) |
| return outs |
|
|
|
|
| def seg_head_forward(seg: SegHead, point_feats: Sequence, frame, *, taps=None, tap_filter=None): |
| """-> ``(labels int64 [N0], probs float32 [N0, 15])`` (torch).""" |
| torch = _torch() |
| tap = _Taps(taps, tap_filter) |
| x = point_feats[4] |
| for s in (3, 2, 1, 0): |
| d = seg.dec[s] |
| up = tap(f"dec{s}.up", gelu(linear(x, d.up))) |
| skip = tap(f"dec{s}.skip", gelu(linear(point_feats[s], d.skip))) |
| x = tap(f"dec{s}.sum", skip + up[torch.from_numpy(frame.pools[s].cluster)]) |
| if d.block is not None: |
| lv = frame.levels[s] |
| x = block_forward(x, d.block, lv.nmap, lv.order[0], lv.inverse[0], tap, cpe_in=skip, name=f"dec{s}.block0") |
| logits = tap("seg.logits", linear(x, seg.classifier)) |
| probs = tap("seg.probs", torch.softmax(logits, dim=1)) |
| labels = probs.argmax(1) |
| tap("seg.labels", labels) |
| return labels, probs |
|
|
|
|
| |
| def conv2d(x, c: Conv2d, residual=None): |
| torch = _torch() |
| y = torch.nn.functional.conv2d(x, _t(c.w), _t(c.b), padding=c.padding, dilation=c.dilation) |
| if residual is not None: |
| y = y + residual |
| return torch.relu(y) if c.relu else y |
|
|
|
|
| def topk_onnx(scores_flat: np.ndarray, k: int) -> np.ndarray: |
| """ONNX ``TopK(largest=1, sorted=1)`` indices with ties broken by the lower index (ORT's order).""" |
| s = np.asarray(scores_flat, np.float32).reshape(-1) |
| return np.argsort(-s, kind="stable")[:k].astype(np.int64) |
|
|
|
|
| def mha(q_in, kv_in, m: MHA, v_in=None): |
| """``[Sq, D]`` x ``[Sk, D]`` -> ``[Sq, D]``: projections, heads, ``softmax((q / div) k^T) v``, out projection.""" |
| torch = _torch() |
| v_in = kv_in if v_in is None else v_in |
| hd = m.q.w.shape[0] // m.heads |
| q = linear(q_in, m.q).reshape(-1, m.heads, hd).transpose(0, 1) / m.scale_div |
| k = linear(kv_in, m.k).reshape(-1, m.heads, hd).transpose(0, 1) |
| v = linear(v_in, m.v).reshape(-1, m.heads, hd).transpose(0, 1) |
| o = torch.empty_like(q) |
| for h in range(m.heads): |
| o[h] = torch.softmax(q[h] @ k[h].T, dim=-1) @ v[h] |
| return linear(o.transpose(0, 1).reshape(q_in.shape[0], -1), m.out) |
|
|
|
|
| def det_head_forward(det: DetHead, point_feats: Sequence, frame, bev_index: np.ndarray, *, |
| topk_indices: Optional[np.ndarray] = None, taps=None, tap_filter=None) -> Dict[str, np.ndarray]: |
| """-> the 9 outputs of ``ptv3_det3d_head.onnx`` (numpy, ONNX shapes and dtypes). ``bev_index``: the host's |
| ``(65536, 9)`` candidate table whose sentinel is the row count of stage 3 (``host.preprocess.bev_index``).""" |
| torch = _torch() |
| tap = _Taps(taps, tap_filter) |
| pf3, pf4 = point_feats[3], point_feats[4] |
| n3 = pf3.shape[0] |
| cluster3 = torch.from_numpy(frame.pools[3].cluster) |
| fused = tap("det.fusion", gelu(linear(pf3, det.fusion_skip)) + gelu(linear(pf4, det.fusion_proj))[cluster3]) |
| proj = tap("det.point_proj", torch.relu(linear(fused, det.point_proj))) |
| table = torch.cat([proj, torch.zeros(1, proj.shape[1], dtype=proj.dtype)], 0) |
| idx = torch.from_numpy(np.where(np.asarray(bev_index) >= n3, n3, bev_index).astype(np.int64)) |
| bev_rows = table[idx].amax(1) |
| tap("det.bev", bev_rows) |
| size = int(round(bev_rows.shape[0] ** 0.5)) |
| x = bev_rows.reshape(1, size, size, -1).permute(0, 3, 1, 2).contiguous() |
| x = tap("det.bev_refine", conv2d(x, det.bev_refine)) |
| x = tap("det.stem", conv2d(x, det.stem)) |
| for i, (c1, c2) in enumerate(det.blocks): |
| x = tap(f"det.block{i}", conv2d(conv2d(x, c1), c2, residual=x)) |
| feat = tap("det.shared", conv2d(x, det.shared)) |
| logits = tap("det.dense_heatmap", conv2d(conv2d(feat, det.heat_hidden), det.heat_out)) |
| heat = tap("det.heat", torch.sigmoid(logits)) |
| local = heat.clone() |
| pc = list(det.pooled_classes) |
| local[:, pc, 1:-1, 1:-1] = torch.nn.functional.max_pool2d(heat[:, pc], 3, stride=1, padding=0) |
| heat_nms = tap("det.heat_nms", heat * (heat == local).to(heat.dtype)) |
| C = heat.shape[1] |
| cells = size * size |
| flat = heat_nms.reshape(1, C, cells) |
| if topk_indices is None: |
| topk_indices = topk_onnx(flat.reshape(-1).numpy(), det.num_proposals) |
| top = np.asarray(topk_indices, np.int64).reshape(-1) |
| tap("det.topk_indices", top) |
| labels = top // cells |
| pos = top % cells |
| pos_t = torch.from_numpy(pos) |
| feat_flat = feat.reshape(1, feat.shape[1], cells) |
| onehot = torch.nn.functional.one_hot(torch.from_numpy(labels), C).to(feat.dtype) |
| query = feat_flat[0][:, pos_t].T + linear(onehot, det.class_encoding) |
| tap("det.query_feat", query) |
| qpos = torch.from_numpy(det.bev_pos)[pos_t] |
| tap("det.query_pos", qpos) |
| qpe = tap("det.query_pos_embed", linear(torch.relu(linear(qpos, det.query_pos[0])), det.query_pos[1])) |
| kpe = tap("det.key_pos_embed", linear(torch.relu(linear(torch.from_numpy(det.bev_pos), det.key_pos[0])), |
| det.key_pos[1])) |
| q = query + qpe |
| q = tap("det.self", q + mha(q, q, det.self_attn)) |
| q = tap("det.norm1", layer_norm(q, det.norm1)) |
| keys = feat_flat[0].T + kpe |
| q = tap("det.cross", q + mha(q, keys, det.cross_attn)) |
| q = tap("det.norm2", layer_norm(q, det.norm2)) |
| q = tap("det.ffn", q + linear(torch.relu(linear(q, det.ffn[0])), det.ffn[1])) |
| q = tap("det.norm3", layer_norm(q, det.norm3)) |
| out: Dict[str, np.ndarray] = {} |
| for name, (h0, h1) in det.heads.items(): |
| y = linear(torch.relu(linear(q, h0)), h1) |
| if name == "center": |
| y = y + qpos |
| tap(f"det.head.{name}", y) |
| out[name] = y.T.unsqueeze(0).numpy().astype(np.float32) |
| out["dense_heatmap"] = logits.numpy().astype(np.float32) |
| out["query_heatmap_score"] = flat[:, :, pos_t].numpy().astype(np.float32) |
| out["query_labels"] = labels.reshape(1, -1).astype(np.int64) |
| return out |
|
|
