| |
| """fp32 PyTorch re-implementation of the deployed TransFusion-L network, module by module, with taps. |
| |
| Built from :mod:`.weights` (the ONNX initializers addressed by node; BN folded where the TT port folds it: the two |
| PFN BatchNorm1d and the deblock2 ConvTranspose BatchNorm), it reproduces ``transfusion.onnx`` (S:transfusion:155-208) |
| from the decorated pillar features on (the in-graph decoration is host pre-processing in this port, computed in |
| float32 exactly as the graph does by ``host.preprocess.decorate_pillars``): |
| |
| - ``pfn``: ``h = relu(x @ W0' + b0')`` (N, 20, 32); ``m = max over ALL 20 slots`` (padded slots included: their rows |
| are ``relu(b0')``, S:transfusion:157); ``y = max over the slots of relu(concat[h, m] @ W1' + b1')`` -> (N, 64); |
| - ``scatter``: ``canvas[0, c, iy, ix] = pfn[p, c]`` on a zero (1, 64, 768, 768) canvas (``ScatterElements``); |
| - ``backbone``: SECOND, 3 blocks of 3x3 convs + ReLU (4 / 6 / 6 layers, first stride 2); |
| - ``neck``: SECONDFPN: conv 2x2 s2, conv 1x1, ConvTranspose 2x2 s2 (BN folded), each + ReLU, concat (1, 384, 192, 192); |
| - ``head`` dense part: shared conv 3x3 384 -> 128 (``lidar_feat``, no activation), heatmap head (3x3 + ReLU, 3x3 -> 5 |
| logits), ``heat = sigmoid``, local max ``heat * (heat == maxpool3x3_s1_p0(heat) with a zero border)`` (border cells |
| always suppressed, every class); |
| - selection: the top ``K = 500`` of the flattened (5 x 36864) local-max heatmap, as ONNX ``TopK`` defines it (largest |
| first, **ties broken by the lower index**: a stable descending sort). ``topk_indices=`` replaces the selection by |
| given indices (teacher forcing, PLAN.md 2.5: the device top-k breaks ties in no defined order, probe P8); |
| - query init: ``class = idx // 36864``, ``pos = idx % 36864``; ``query = lidar_feat[:, pos] + class_table[class]`` |
| (the Conv1d of a one-hot is the table row ``W[:, c] + b``); ``query_pos = (pos % 192 + 0.5, pos // 192 + 0.5)``, |
| computed arithmetically in float32 (bit-equal to the graph's ``bev_pos`` gather, checked at load); |
| - decoder (one layer, post-norm, 8 heads x 16): self-attention on ``x + pe(query_pos)`` (Q = K = V), cross-attention |
| of ``x + pe(query_pos)`` over ``lidar_feat^T + pe(bev_pos)`` (36,864 keys), FFN 128 -> 256 -> 128, LayerNorm eps |
| 1e-5 after each; ``Q`` is divided by ``sqrt(16) = 4`` before ``Q K^T`` (the graph's ``Div``); |
| - prediction heads (Conv1d 128 -> 64 + ReLU, Conv1d 64 -> C for center 2, height 1, dim 3, rot 2, vel 2, heatmap 5); |
| - outputs: ``bbox_pred0 = [center + query_pos (float32), height, dim, vel]`` (1, 8, K), ``dir_cls_pred0 = rot`` |
| (1, 2, K), ``cls_score0 = query_heat * sigmoid(heads.heatmap)`` (1, 5, K) with ``query_heat`` the local-max |
| heatmap of all 5 classes at the query's cell (no one-hot: mmdeploy.patch). |
| |
| The TT-friendly equivalent forms the port uses (PFN split, K-split shared conv, constant cross-attention K / V tables, |
| merged prediction heads, the masked local-max chain, 32-slot pillar tiles, head-dim padding) are in :mod:`.rewrites` |
| and proven equal to these modules by ``tests/test_reference_rewrites_host.py``. |
| |
| Taps (``ttaw.golden.TapRegistry``; stable names, the keys of the device gates; batch axis kept as the ONNX tensors |
| have it): :data:`MODULE_TAPS` plus the opt-in per-layer / large taps :data:`EXTRA_TAPS`. Values are float32 numpy |
| (``head.topk_indices`` int64). torch is imported lazily; the module runs where ttnn is absent. |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any, Dict, Optional |
|
|
| import numpy as np |
|
|
| from ..ttaw.golden import NULL_TAPS, TapRegistry |
| from .weights import TransFusionWeights |
|
|
| __all__ = ["TransFusionTorch", "MODULE_TAPS", "EXTRA_TAPS", "SELECTION_TAPS", "OUTPUT_NAMES", "stable_topk"] |
|
|
| OUTPUT_NAMES = ("cls_score0", "bbox_pred0", "dir_cls_pred0") |
| |
| MODULE_TAPS = ("pfn.input", "pfn", "canvas", "backbone.block0", "backbone.block1", "backbone.block2", |
| "neck.deblock0", "neck.deblock1", "neck.deblock2", "head.shared", "head.heatmap", "head.heat_nms", |
| "head.topk_values", "head.topk_indices", "query.feat", "query.pos", "query.heat", |
| "decoder.query_pos_embed", "decoder.self_attn", "decoder.cross_attn", "decoder.out", "heads.center", |
| "heads.height", "heads.dim", "heads.rot", "heads.vel", "heads.heatmap") + OUTPUT_NAMES |
| |
| EXTRA_TAPS = ("pfn.layer0", "neck", "head.heatmap.hidden", "decoder.key_pos_embed") + tuple( |
| f"backbone.block{b}.conv{j}" for b, n in ((0, 4), (1, 6), (2, 6)) for j in range(n)) |
| |
| SELECTION_TAPS = ("head.topk_values", "head.topk_indices", "query.feat", "query.pos", "query.heat", |
| "decoder.query_pos_embed", "decoder.self_attn", "decoder.cross_attn", "decoder.out", |
| "heads.center", "heads.height", "heads.dim", "heads.rot", "heads.vel", "heads.heatmap") + OUTPUT_NAMES |
|
|
|
|
| def _torch(): |
| import torch |
|
|
| return torch |
|
|
|
|
| def _tap(taps: TapRegistry, name: str, value): |
| """Record a torch tensor as a numpy copy (float32, or int64 for integer tensors) and return it unchanged; never |
| through ``ttaw.tensors.to_numpy`` (which would import ttnn).""" |
| if taps.enabled and taps.wants(name): |
| a = value.detach().cpu().numpy() |
| taps.tap(name, a.astype(np.int64 if a.dtype.kind in "iu" else np.float32, copy=True)) |
| return value |
|
|
|
|
| def stable_topk(scores, k: int): |
| """ONNX ``TopK(largest=1, sorted=1)`` on the last axis: values descending, equal values in ascending index order |
| (a stable descending sort). Returns ``(values, indices)``.""" |
| torch = _torch() |
| order = torch.sort(scores, dim=-1, descending=True, stable=True) |
| return order.values[..., :k], order.indices[..., :k] |
|
|
|
|
| class TransFusionTorch: |
| """The fp32 reference network. ``threads`` caps torch's intra-op threads (shared host: <= 4).""" |
|
|
| def __init__(self, weights: TransFusionWeights, *, threads: Optional[int] = None, device: str = "cpu"): |
| torch = _torch() |
| if threads is not None: |
| torch.set_num_threads(int(threads)) |
| self.weights = weights |
| self.device = device |
| t = self._t |
| p = weights.pfn |
| self.w0, self.b0, self.w1, self.b1 = t(p.w0), t(p.b0), t(p.w1), t(p.b1) |
| self._layers: Dict[str, tuple] = {} |
| for layer in weights.conv_layers(): |
| self._layers[layer.name] = (layer, t(layer.weight), t(layer.bias)) |
| h = weights.head |
| d = h.decoder |
| for layer in (*d.self_posembed, *d.cross_posembed) + tuple(x for _, a, b in h.prediction_heads for x in (a, b)): |
| self._layers[layer.name] = (layer, t(layer.weight), t(layer.bias)) |
| self.class_table = t(h.class_table) |
| self.bev_pos = t(h.bev_pos) |
| self.K = int(h.num_proposals) |
| self.grid_w, self.grid_h = weights.grid |
| self.feat_w = int(round(np.sqrt(h.bev_pos.shape[0]))) |
| self._attn = {"self_attn": self._attn_tensors(d.self_attn), "cross_attn": self._attn_tensors(d.cross_attn)} |
| self._norms = [(t(g), t(b), float(eps)) for g, b, eps in d.norms] |
| self._ffn = (t(d.ffn_w1), t(d.ffn_b1), t(d.ffn_w2), t(d.ffn_b2)) |
| self._kpe = None |
|
|
| def _t(self, a: Any): |
| torch = _torch() |
| return torch.from_numpy(np.array(a, dtype=np.float32, copy=True)).to(self.device) |
|
|
| def _attn_tensors(self, a) -> dict: |
| t = self._t |
| return {"wq": t(a.wq), "bq": t(a.bq), "wk": t(a.wk), "bk": t(a.bk), "wv": t(a.wv), "bv": t(a.bv), |
| "wo": t(a.wo), "bo": t(a.bo), "heads": a.num_heads, "hd": a.head_dim, "div": float(a.q_divisor)} |
|
|
| |
| def pfn(self, features: Any, taps: TapRegistry = NULL_TAPS): |
| """(N, 20, 11) decorated pillar features (padded slots zero) -> (N, 64) pillar features.""" |
| torch = _torch() |
| x = _tap(taps, "pfn.input", torch.as_tensor(np.asarray(features, dtype=np.float32), device=self.device)) |
| h = _tap(taps, "pfn.layer0", torch.relu(torch.matmul(x, self.w0) + self.b0)) |
| m = h.amax(dim=1, keepdim=True) |
| y = torch.relu(torch.matmul(torch.cat([h, m.expand_as(h)], dim=2), self.w1) + self.b1) |
| return _tap(taps, "pfn", y.amax(dim=1)) |
|
|
| def scatter(self, pillar_features: Any, coords: np.ndarray, taps: TapRegistry = NULL_TAPS): |
| """``ScatterElements`` into the zero canvas: (N, 64) -> (1, 64, 768, 768) with ``canvas[c, iy, ix]``. |
| ``coords``: (N, >=2) with (.., iy, ix) in the last two columns (``(0, iy, ix)`` or ONNX ``coors``).""" |
| torch = _torch() |
| pf = torch.as_tensor(pillar_features, device=self.device) |
| canvas = torch.zeros((1, pf.shape[1], self.grid_h, self.grid_w), dtype=torch.float32, device=self.device) |
| c = torch.as_tensor(np.asarray(coords)[:, -2:].astype(np.int64), device=self.device) |
| canvas[0, :, c[:, 0], c[:, 1]] = pf.T |
| return _tap(taps, "canvas", canvas) |
|
|
| def layer(self, name: str, x): |
| torch = _torch() |
| F = torch.nn.functional |
| layer, w, b = self._layers[name] |
| if layer.kind == "conv": |
| y = F.conv2d(x, w, b, stride=layer.stride, padding=layer.padding) |
| elif layer.kind == "conv_transpose": |
| y = F.conv_transpose2d(x, w, b, stride=layer.stride, padding=layer.padding) |
| else: |
| y = F.conv1d(x, w, b, stride=layer.stride, padding=layer.padding) |
| return torch.relu(y) if layer.relu else y |
|
|
| def backbone(self, canvas, taps: TapRegistry = NULL_TAPS): |
| outs, x = [], canvas |
| for i, block in enumerate(self.weights.blocks): |
| for layer in block: |
| x = _tap(taps, layer.name, self.layer(layer.name, x)) |
| outs.append(_tap(taps, f"backbone.block{i}", x)) |
| return outs |
|
|
| def neck(self, blocks, taps: TapRegistry = NULL_TAPS): |
| torch = _torch() |
| ups = [_tap(taps, layer.name, self.layer(layer.name, x)) for layer, x in zip(self.weights.deblocks, blocks)] |
| return _tap(taps, "neck", torch.cat(ups, dim=1)) |
|
|
| def dense_head(self, neck, taps: TapRegistry = NULL_TAPS): |
| """-> (lidar_feat (1, 128, 192, 192), heat_nms (1, 5, 192, 192)).""" |
| torch = _torch() |
| F = torch.nn.functional |
| lidar = _tap(taps, "head.shared", self.layer("head.shared", neck)) |
| hid = _tap(taps, "head.heatmap.hidden", self.layer("head.heatmap.hidden", lidar)) |
| logits = _tap(taps, "head.heatmap", self.layer("head.heatmap", hid)) |
| heat = torch.sigmoid(logits) |
| k = int(self.weights.head.nms_kernel) |
| pad = k // 2 |
| local_max = torch.zeros_like(heat) |
| local_max[:, :, pad:-pad, pad:-pad] = F.max_pool2d(heat, k, 1, 0) |
| heat_nms = _tap(taps, "head.heat_nms", heat * (heat == local_max).to(heat.dtype)) |
| return lidar, heat_nms |
|
|
| def select(self, heat_nms, topk_indices: Optional[Any] = None, taps: TapRegistry = NULL_TAPS): |
| """Top-K of the flattened local-max heatmap -> (values (1, K), indices (1, K) int64); given ``topk_indices`` |
| (teacher forcing) are used as they are and their heat values gathered.""" |
| torch = _torch() |
| flat = heat_nms.reshape(1, -1) |
| if topk_indices is None: |
| values, idx = stable_topk(flat, self.K) |
| else: |
| idx = torch.as_tensor(np.asarray(topk_indices, dtype=np.int64).reshape(1, -1), device=self.device) |
| if idx.shape[1] != self.K or int(idx.min()) < 0 or int(idx.max()) >= flat.shape[1]: |
| raise ValueError(f"topk_indices must be {self.K} indices in [0, {flat.shape[1]})") |
| values = flat.gather(1, idx) |
| return _tap(taps, "head.topk_values", values), _tap(taps, "head.topk_indices", idx) |
|
|
| def query_init(self, lidar, heat_nms, idx, taps: TapRegistry = NULL_TAPS): |
| """-> (query (1, 128, K), query_pos (1, K, 2), query_heat (1, 5, K)).""" |
| torch = _torch() |
| n_cells = lidar.shape[2] * lidar.shape[3] |
| cls_id = torch.div(idx, n_cells, rounding_mode="floor") |
| pos = idx - cls_id * n_cells |
| flat = lidar.reshape(1, lidar.shape[1], -1) |
| q = flat.gather(2, pos[:, None, :].expand(-1, flat.shape[1], -1)) + self.class_table[cls_id[0]].T[None] |
| w = self.feat_w |
| qpos = torch.stack([(pos % w).to(torch.float32) + 0.5, |
| torch.div(pos, w, rounding_mode="floor").to(torch.float32) + 0.5], dim=-1) |
| heat_flat = heat_nms.reshape(1, heat_nms.shape[1], -1) |
| qheat = heat_flat.gather(2, pos[:, None, :].expand(-1, heat_flat.shape[1], -1)) |
| return _tap(taps, "query.feat", q), _tap(taps, "query.pos", qpos), _tap(taps, "query.heat", qheat) |
|
|
| def _posembed(self, which: str, pos): |
| """(P, 2) -> (P, 128): Conv1d 2 -> 128 + ReLU, Conv1d 128 -> 128 (as 1x1 convs over the P positions).""" |
| x = pos.T[None] |
| x = self.layer(f"decoder.{which}.0", x) |
| return self.layer(f"decoder.{which}.1", x)[0].T |
|
|
| def key_pos_embed(self, taps: TapRegistry = NULL_TAPS): |
| """``cross_posembed(bev_pos)`` (36864, 128): input-independent, computed once.""" |
| torch = _torch() |
| if self._kpe is None: |
| with torch.no_grad(): |
| self._kpe = self._posembed("cross_posembed", self.bev_pos) |
| return _tap(taps, "decoder.key_pos_embed", self._kpe[None])[0] |
|
|
| def _mha(self, which: str, q_in, kv_in, v_in=None): |
| """Multi-head attention as exported: projections, per-head ``softmax((q / div) @ k^T) @ v``, out proj.""" |
| torch = _torch() |
| a = self._attn[which] |
| v_in = kv_in if v_in is None else v_in |
| q = q_in @ a["wq"] + a["bq"] |
| k = kv_in @ a["wk"] + a["bk"] |
| v = v_in @ a["wv"] + a["bv"] |
| H, D = a["heads"], a["hd"] |
| lq, lk = q.shape[0], k.shape[0] |
| q = q.reshape(lq, H, D).transpose(0, 1) / a["div"] |
| k = k.reshape(lk, H, D).permute(1, 2, 0) |
| v = v.reshape(lk, H, D).transpose(0, 1) |
| o = torch.softmax(q @ k, dim=-1) @ v |
| o = o.transpose(0, 1).reshape(lq, H * D) |
| return o @ a["wo"] + a["bo"] |
|
|
| def _ln(self, i: int, x): |
| torch = _torch() |
| g, b, eps = self._norms[i] |
| return torch.nn.functional.layer_norm(x, (x.shape[-1],), g, b, eps) |
|
|
| def decoder(self, query, qpos, lidar, taps: TapRegistry = NULL_TAPS): |
| """One post-norm decoder layer -> (K, 128).""" |
| torch = _torch() |
| qpe = _tap(taps, "decoder.query_pos_embed", self._posembed("self_posembed", qpos[0])[None])[0] |
| kpe = self.key_pos_embed(taps) |
| x = query[0].T |
| key = lidar.reshape(lidar.shape[1], -1).T + kpe |
| s = x + qpe |
| x = _tap(taps, "decoder.self_attn", self._ln(0, x + self._mha("self_attn", s, s))[None])[0] |
| x = _tap(taps, "decoder.cross_attn", self._ln(1, x + self._mha("cross_attn", x + qpe, key))[None])[0] |
| w1, b1, w2, b2 = self._ffn |
| ff = torch.relu(x @ w1 + b1) @ w2 + b2 |
| return _tap(taps, "decoder.out", self._ln(2, x + ff)[None])[0] |
|
|
| def prediction_heads(self, x, qpos, qheat, taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]: |
| torch = _torch() |
| y = x.T[None] |
| out = {} |
| for name, hid, fin in self.weights.head.prediction_heads: |
| out[name] = _tap(taps, fin.name, self.layer(fin.name, self.layer(hid.name, y))) |
| center = out["center"] + qpos.transpose(1, 2) |
| cls = qheat * torch.sigmoid(out["heatmap"]) |
| box = torch.cat([center, out["height"], out["dim"], out["vel"]], dim=1) |
| return {"cls_score0": _tap(taps, "cls_score0", cls), "bbox_pred0": _tap(taps, "bbox_pred0", box), |
| "dir_cls_pred0": _tap(taps, "dir_cls_pred0", out["rot"])} |
|
|
| |
| def forward(self, features: np.ndarray, coords: np.ndarray, num_pillars: int, *, |
| topk_indices: Optional[Any] = None, taps: TapRegistry = NULL_TAPS) -> Dict[str, np.ndarray]: |
| """Decorated pillars (the first ``num_pillars`` rows of ``features`` / ``coords`` are used) -> the three |
| graph outputs as float32 numpy: ``cls_score0`` (1, 5, K), ``bbox_pred0`` (1, 8, K), ``dir_cls_pred0`` |
| (1, 2, K). ``topk_indices`` (K,) or (1, K): teacher-forced query selection.""" |
| torch = _torch() |
| n = int(num_pillars) |
| with torch.no_grad(): |
| pf = self.pfn(np.asarray(features)[:n], taps) |
| canvas = self.scatter(pf, np.asarray(coords)[:n], taps) |
| neck = self.neck(self.backbone(canvas, taps), taps) |
| lidar, heat_nms = self.dense_head(neck, taps) |
| _, idx = self.select(heat_nms, topk_indices, taps) |
| query, qpos, qheat = self.query_init(lidar, heat_nms, idx, taps) |
| x = self.decoder(query, qpos, lidar, taps) |
| out = self.prediction_heads(x, qpos, qheat, taps) |
| return {k: v.cpu().numpy().astype(np.float32) for k, v in out.items()} |
|
|
| def describe(self) -> Dict[str, Any]: |
| w = self.weights |
| d = w.head.decoder |
| return {"onnx_sha256": w.sha256, "grid": list(w.grid), "num_proposals": self.K, |
| "blocks": [[(lay.in_channels, lay.out_channels, lay.stride) for lay in b] for b in w.blocks], |
| "deblocks": [(lay.kind, lay.in_channels, lay.out_channels, lay.kernel, lay.stride) for lay in w.deblocks], |
| "decoder": {"heads": d.self_attn.num_heads, "head_dim": d.self_attn.head_dim, |
| "ffn": int(d.ffn_w1.shape[1])}, |
| "prediction_heads": {n: fin.out_channels for n, _, fin in w.head.prediction_heads}} |
|
|
|
|
| def layer_names(weights: TransFusionWeights) -> list: |
| """Every per-layer conv tap name (backbone convs, deblocks, head convs).""" |
| return [layer.name for layer in weights.conv_layers()] |
|
|
|
|
