# SPDX-License-Identifier: Apache-2.0 """fp32 torch rebuild of the deployed FRNet graph (AWML ``frustum_encoder.py:76-137``, ``frnet_backbone.py:252-356``, ``frnet_head.py:52-70`` in deploy mode), with per-module taps. Ported from the verified research reference ``research/frnet/scripts/frnet_torch_ref.py`` (max |diff| vs ORT 1.6e-5 .. 3.0e-5 on the public goldens), on the weights as the TT port holds them (``reference/weights.py``: point-layer BatchNorms folded, ``pre_norm`` as x * scale + shift). Input: the four ONNX inputs of one frame, **unpadded** (``host/preprocess.py`` ``NetworkInputs``); the graph sizes its voxel tensors as ``max(inverse_map) + 1`` (S:frnet:255-262). Every sparse op is deterministic and follows ORT: - cell mean: index-add of the points / of ones, divided (``ScatterElements(add)`` x2 + ``Div``); - per-cell max (``ScatterElements(max)``, init ``-FLT_MAX``): ``scatter_reduce(amax, include_self=True)``; - frustum2pixel (``ScatterND`` into zeros, NHWC -> NCHW): at strides 2 / 4 / 8 several cells fall on one pixel and ORT applies the updates sequentially, so the **last voxel row** wins (S:frnet:285-287); - pixel2point (``Gather``): row ``(row // s) * (W // s) + col // s`` of the NHWC map; - decoder: bilinear ``align_corners=True`` resize of the five maps to H x W, channel concat. Taps (``ttaw.golden.TapRegistry``; names in ``weights.MODULE_TAPS`` / ``OPT_IN_TAPS``): point tensors ``(N, C)``, voxel tensors ``(M, C)``, maps ``(1, C, h, w)`` NCHW, all float32 numpy. ``forward`` returns ``pred_probs`` ``(N, 27)`` float32. torch is imported here (research venv / tt-env), never by ``api.py`` at module level. """ from __future__ import annotations from typing import Any, Optional import numpy as np import torch import torch.nn.functional as F from ..ttaw.golden import NULL_TAPS, TapRegistry from .weights import Conv, FRNetWeights, Linear __all__ = ["FRNetTorch", "scatter_max", "frustum2pixel", "pixel2point", "TorchLayers"] def _t(a: np.ndarray) -> torch.Tensor: a = np.ascontiguousarray(a) return torch.from_numpy(a if a.flags.writeable else a.copy()) class TorchLayers: """The weights as torch tensors (shared by the literal model and the rewrites).""" def __init__(self, weights: FRNetWeights): self.w = weights self.lin = {lin.name: (_t(lin.weight), _t(lin.bias), lin.relu) for lin in weights.linears} self.conv = {c.name: (_t(c.weight), _t(c.bias), c.stride, c.padding, c.activation) for c in weights.convs} self.pre_scale = _t(weights.pre_norm_scale) self.pre_shift = _t(weights.pre_norm_shift) def linear(self, x: torch.Tensor, layer: Linear) -> torch.Tensor: w, b, relu = self.lin[layer.name] y = torch.addmm(b, x, w) return F.relu(y) if relu else y def conv2d(self, x: torch.Tensor, layer: Conv, *, activation: bool = True) -> torch.Tensor: w, b, stride, pad, act = self.conv[layer.name] y = F.conv2d(x, w, b, stride=stride, padding=pad) if activation and act == "hardswish": y = F.hardswish(y) elif activation and act == "sigmoid": y = torch.sigmoid(y) return y def scatter_max(src: torch.Tensor, inverse_map: torch.Tensor, m: int) -> torch.Tensor: """``ScatterElements(reduction=max)`` into ``-FLT_MAX`` rows (every row of a frame owns >= 1 point).""" out = torch.full((m, src.shape[1]), torch.finfo(src.dtype).min, dtype=src.dtype) return out.scatter_reduce(0, inverse_map.unsqueeze(-1).expand(-1, src.shape[1]), src, reduce="amax", include_self=True) def frustum2pixel(feats: torch.Tensor, voxel_coors: np.ndarray, hw: tuple, stride: int) -> torch.Tensor: """``ScatterND`` of voxel rows into a zero ``(1, h, w, C)`` image at ``voxel_coors // stride`` (last voxel wins on duplicates), returned NCHW.""" H, W = hw h, w = H // stride, W // stride lin = (voxel_coors[:, 1] // stride) * w + voxel_coors[:, 2] // stride last = np.full(h * w, -1, np.int64) np.maximum.at(last, lin, np.arange(len(lin), dtype=np.int64)) pix = np.flatnonzero(last >= 0) img = torch.zeros(h * w, feats.shape[1], dtype=feats.dtype) img[_t(pix)] = feats[_t(last[pix])] return img.reshape(1, h, w, -1).permute(0, 3, 1, 2).contiguous() def pixel2point(x: torch.Tensor, coors: np.ndarray, stride: int) -> torch.Tensor: """``Gather`` of the NHWC rows of ``x`` (NCHW) at each point's cell ``// stride``.""" _, c, h, w = x.shape flat = x.permute(0, 2, 3, 1).reshape(h * w, c) return flat[_t((coors[:, 1] // stride) * w + coors[:, 2] // stride)] class FRNetTorch: """``FRNetTorch(weights)(points, coors, voxel_coors, inverse_map, taps=TapRegistry())`` -> pred_probs.""" def __init__(self, weights: FRNetWeights, *, threads: Optional[int] = None): if threads: torch.set_num_threads(int(threads)) self.weights = weights self.layers = TorchLayers(weights) def __call__(self, *args: Any, **kwargs: Any) -> np.ndarray: return self.forward(*args, **kwargs) @torch.no_grad() def forward(self, points: np.ndarray, coors: np.ndarray, voxel_coors: np.ndarray, inverse_map: np.ndarray, taps: TapRegistry = NULL_TAPS) -> np.ndarray: W, L = self.weights, self.layers hw = W.frustum_hw m = int(inverse_map.max()) + 1 voxel_coors = np.asarray(voxel_coors)[:m] coors = np.asarray(coors) pts = _t(np.asarray(points, np.float32)) inv = _t(np.asarray(inverse_map, np.int64)) def tap(name: str, value: torch.Tensor) -> torch.Tensor: if taps.wants(name): taps.tap(name, value.numpy()) return value # ---- FrustumFeatureEncoder dist = torch.linalg.norm(pts[:, :3], dim=1, keepdim=True) vsum = torch.zeros(m, 4).index_add_(0, inv, pts) vcnt = torch.zeros(m, 4).index_add_(0, inv, torch.ones_like(pts)) f_cluster = pts[:, :3] - (vsum / vcnt)[inv][:, :3] x = torch.cat([pts, dist, f_cluster], 1) * L.pre_scale + L.pre_shift tap("encoder.input", x) pf = [] for i, lin in enumerate(W.ffe): x = L.linear(x, lin) pf.append(tap(f"encoder.pf{i}", x)) vmax = tap("encoder.voxel_max", scatter_max(x, inv, m)) vfeat = tap("encoder.voxel_feats", L.linear(vmax, W.compression)) # ---- backbone: stem, point stem, fusion stem img = tap("backbone.frustum", frustum2pixel(vfeat, voxel_coors, hw, 1)) for c in W.stem: img = L.conv2d(img, c) tap("backbone.stem", img) p = L.linear(torch.cat([pixel2point(img, coors, 1), pf[-1]], 1), W.point_stem) tap("backbone.point_stem", p) seg = tap("backbone.segmax0", scatter_max(p, inv, m)) x = L.conv2d(torch.cat([frustum2pixel(seg, voxel_coors, hw, 1), img], 1), W.fusion_stem) tap("backbone.fusion_stem", x) maps, pts_out = [x], [p] # ---- four stages: BasicBlocks, point fusion, pixel fusion, attention for st in W.stages: li, s = st.index, st.stride for bi, (c1, c2, ds) in enumerate(st.blocks): y = L.conv2d(L.conv2d(x, c1), c2) idt = L.conv2d(x, ds) if ds is not None else x x = tap(f"backbone.layer{li}.block{bi}", F.hardswish(y + idt)) tap(f"backbone.layer{li}", x) p = L.linear(torch.cat([pixel2point(x, coors, s), p], 1), st.point_fusion) tap(f"backbone.point_fusion{li}", p) seg = tap(f"backbone.segmax{li}", scatter_max(p, inv, m)) fuse = L.conv2d(torch.cat([frustum2pixel(seg, voxel_coors, hw, s), x], 1), st.pixel_fusion) tap(f"backbone.pixel_fusion{li}", fuse) att = tap(f"backbone.attention{li}", L.conv2d(L.conv2d(fuse, st.attention[0]), st.attention[1])) x = tap(f"backbone.stage{li}", fuse * att + x) maps.append(x) pts_out.append(p) # ---- decoder (align-corners bilinear up-sampling to H x W, fuse convs) and the point branch ups = [tap(f"decoder.up{k}", F.interpolate(mp, size=hw, mode="bilinear", align_corners=True)) for k, mp in enumerate(maps)] y = tap("decoder.fuse1", L.conv2d(torch.cat(ups, 1), W.fuse_layer1)) y = tap("decoder.fuse2", L.conv2d(y, W.fuse_layer2)) q = tap("decoder.point_fuse1", L.linear(torch.cat(pts_out, 1), W.point_fuse[0])) q = tap("decoder.point_fuse2", L.linear(q, W.point_fuse[1])) # ---- FRHead h = pixel2point(y, coors, 1) for i, (lin, skip) in enumerate(zip(W.head_mlps, (q, pf[2], pf[1], pf[0]))): h = tap(f"head.mlp{i}", L.linear(h, lin) + skip) logits = tap("head.logits", L.linear(h, W.conv_seg)) return tap("pred_probs", torch.softmax(logits, -1)).numpy()