| |
| """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 |
|
|
| |
| 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)) |
|
|
| |
| 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] |
|
|
| |
| 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) |
|
|
| |
| 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])) |
|
|
| |
| 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() |
|
|
