File size: 4,522 Bytes
77e3c67 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | # SPDX-License-Identifier: Apache-2.0
"""The three in-network TensorRT plugins of ``bevdet_one_lt_d.onnx`` in torch fp32 (the vendor fp32 path).
Clean-room from research/bevdet/SPEC.md section 4.2 (no vendor code), ported from the verified research reference
(``research/bevdet/scripts/bevdet_common.py``):
- :func:`bev_pool`: ``bevdet::BEVPool``: ``out[c, ranks_bev[k]] += feat[ranks_feat[k] -> (n, c, hw)] *
depth_flat[ranks_depth[k]]`` over every table entry, zero-initialised (the intended semantics; the TRT plugin
leaves never-hit cells uninitialised, S:162). One ``index_add_`` in fp32.
- :func:`align_bev`: ``bevdet::AlignBEV``: for history slot n and BEV pixel (h, w): ``ix = a w + b h + c``,
``iy = d w + e h + f`` (pixel coordinates), bilinear over the 4 neighbours with zero outside
(= ``grid_sample(align_corners=True, padding_mode="zeros")``, as in training). The 4 taps are evaluated literally
in float32 like the CUDA kernel, so an identity transform is an exact copy.
- :func:`gather_bev`: ``bevdet::GatherBEV``: ``cat(curr, s_1 .. s_8)`` with ``s_j = aligned[j - 1]`` if ``flag``
else ``curr``; channel order [t, t-1, ..., t-8].
Inputs and outputs are torch tensors or numpy arrays (outputs: torch). The Preprocess plugin is host work
(``host.preprocess``).
"""
from __future__ import annotations
from typing import Any
import numpy as np
from ..host.calib import BEVPoolTables
from .config import BEVDetConfig
__all__ = ["bev_pool", "align_bev", "gather_bev"]
def _t(x: Any):
import torch
return torch.as_tensor(np.asarray(x, np.float32)) if not isinstance(x, torch.Tensor) else x.to(torch.float32)
def bev_pool(depth: Any, feat: Any, tables: BEVPoolTables, cfg: BEVDetConfig):
"""depth [N, D, H, W] (softmax), feat [N, C, H, W] -> curr [1, C, ny, nx] (fp32, cells never hit are 0)."""
import torch
depth_t, feat_t = _t(depth), _t(feat)
N, C, H, W = feat_t.shape
ny, nx = cfg.bev_hw
nz = cfg.bev_xyz[2]
dflat = depth_t.reshape(-1)
fflat = feat_t.reshape(N, C, H * W).permute(1, 0, 2).reshape(C, N * H * W)
rd = torch.from_numpy(tables.ranks_depth.astype(np.int64))
rf = torch.from_numpy(tables.ranks_feat.astype(np.int64))
rb = torch.from_numpy(tables.ranks_bev.astype(np.int64))
contrib = fflat[:, rf] * dflat[rd][None, :]
out = torch.zeros(C, nx * ny * nz, dtype=torch.float32)
out.index_add_(1, rb, contrib)
return out.reshape(1, C, ny, nx)
def align_bev(adj_feats: Any, transforms: Any):
"""adj_feats [1, A, C, H, W], transforms [1, A, 6] -> aligned [1, A, C, H, W] (bilinear, zero outside)."""
import torch
adj = _t(adj_feats)
trs = _t(transforms)
B, A, C, H, W = adj.shape
out = torch.zeros_like(adj)
hh, ww = torch.meshgrid(torch.arange(H, dtype=torch.float32), torch.arange(W, dtype=torch.float32),
indexing="ij")
identity = torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float32)
for n in range(A):
if torch.equal(trs[0, n], identity): # one tap with weight 1, three with weight 0: a copy
out[0, n] = adj[0, n]
continue
a, b, c0, d, e, f = (trs[0, n, k] for k in range(6))
ix = a * ww + b * hh + c0 # float32, like the kernel
iy = d * ww + e * hh + f
x0 = torch.floor(ix).double()
y0 = torch.floor(iy).double()
x1, y1 = x0 + 1, y0 + 1
ixd, iyd = ix.double(), iy.double()
# weights: exact differences (float64), one rounding of the product to float32 (research reference)
taps = ((y0, x0, (x1 - ixd) * (y1 - iyd)), (y0, x1, (ixd - x0) * (y1 - iyd)),
(y1, x0, (x1 - ixd) * (iyd - y0)), (y1, x1, (ixd - x0) * (iyd - y0)))
src = adj[0, n]
acc = torch.zeros((C, H, W), dtype=torch.float32)
for yy, xx, wt in taps:
ok = (yy >= 0) & (yy < H) & (xx >= 0) & (xx < W)
yi = yy.clamp(0, H - 1).long()
xi = xx.clamp(0, W - 1).long()
vals = src[:, yi, xi]
acc += torch.where(ok[None], vals * wt.float()[None], torch.zeros((), dtype=torch.float32))
out[0, n] = acc
return out
def gather_bev(aligned: Any, curr: Any, flag: int):
"""aligned [1, A, C, H, W], curr [1, C, H, W] -> [1, (1 + A) C, H, W]."""
import torch
al, cu = _t(aligned), _t(curr)
A = al.shape[1]
parts = [cu[0]] + [cu[0] if int(flag) == 0 else al[0, n] for n in range(A)]
return torch.cat(parts, 0)[None]
|