| |
| """Pure-PyTorch fp32 re-implementation of ``meteor_v157c3Z.onnx`` with per-module taps. |
| |
| A node-by-node transliteration of the deployed graph (2,021 ONNX nodes; ``research/meteor/SPEC.md`` sections 4.2-4.4, |
| cited ``S:<line>``), written from the graph itself (no upstream code is imported or executed): every convolution runs |
| with the weights and the attributes of its ONNX node (``reference.weights``), every non-conv constant (attention |
| pools, LayerNorm, gates, pooling matrices, lift ground points) is read from its consuming node, and the arithmetic |
| keeps the graph's operation order where it decides a rounding or a discrete choice (lift projection, validity and |
| depth-bin selection, risk sampling, decoupled decoder, depth_mean). Shape glue (Shape / Gather / Concat / Reshape of |
| shapes, ~470 nodes) disappears. ``MeteorNet.forward`` returns the 19 graph outputs with the ONNX names, dtypes and |
| shapes; ``taps=`` records the intermediates below (names shared with ``reference.ort.ORT_TAPS``, the device gates |
| and the goldens). |
| |
| Module boundaries (what a TT stage computes) and taps: |
| |
| ========================== =========================== ==================================================== |
| tap shape graph tensor / meaning |
| ========================== =========================== ==================================================== |
| ``img.input`` [8, 3, 432, 768] uint8 / 255 (+ ImageNet mean/std with input_norm=imagenet) |
| ``img.stem`` [8, 64, 108, 192] 7x7 s2 + ReLU + maxpool 3x3 s2 |
| ``img.layer1..4`` [8, 64..512, ...] ResNet-34 stages (BasicBlocks, BN folded) |
| ``img.fpn`` [8, 160, 108, 192] lat1..4 + bilinear (x2, x4, 14x24 -> 108x192) + fuse 3x3 ReLU |
| ``seg2d.logits`` [8, 21, 108, 192] SegHeadED raw logits (the PointPainting input, unclamped) |
| ``seg2d.clamped`` [8, 21, 108, 192] nan_to_num + clamp +-30 (the seg2d argmax input only) |
| ``depth.logits`` / ``.prob`` [8, 64, 108, 192] depth head 1x1 / softmax over the 64 linear bins |
| ``ctx.painted`` [8, 96, 108, 192] ctx 1x1 + paint_proj(softmax(seg)[2,3,4,5,6,7,8,13]) |
| ``lift.grid`` [8, 100000, 1, 2] normalised (u, v) of the 400x250 ground points per camera |
| ``lift.valid`` [8, 100000] bool z > 0.5, inside the image, range < 90 m |
| ``lift.weight`` [8, 1, 100000] (interpolated bin probability + 0.05) * valid |
| ``lift.bev`` [1, 96, 400, 250] sum_n ctx_s w / max(sum_n w, 1e-4) |
| ``bev.raw`` [1, 96, 800, 500] bilinear x2 of lift.bev |
| ``bev.gate`` [1, 1, 800, 500] 4 softmax(tgate(raw))[0] |
| ``bev.fused`` [1, 96, 800, 500] raw + tfuse(raw * gate) |
| ``lane.pre`` [1, 9, 800, 500] lane decoder + lane_branch on channels 4..6 |
| ``det.feat`` [1, 128, 400, 250] det stem (ED) output |
| ``det.hm_pre/.reg_pre`` [1, 2 | 6, 400, 250] heads before the box refiner |
| ``occ.feat`` [1, 192, 200, 200] occ stem on raw rows 200:600, cols 50:450 |
| ``traj.stem / .feat`` [1, 128 | 256, 400, 250] traj stem; cat(stem, det.feat) + agent_delta |
| ``traj.agent_delta`` [1, 256, 1, 1] agent attention pool -> 1x1 conv |
| ``stat.feat`` [1, 48, 400, 250] delta_stat convs on avgpool2(|raw|) |
| ``risk.feat`` [1, 64, 400, 250] risk ConvBlocks on fused rows 200:600, cols 125:375 |
| ``ego.*`` [1, 42] (and others) planner chain: stem, mlp, attn, plus_attn, sem_in, sem, |
| plus_sem, kin, plus_kin, risk_sample, after_risk, |
| dec_head, dec_wp, after_dec |
| ``bev.fused_mean`` [1, 96] global mean of the fused BEV (computed once; S:368-369) |
| ``refiner.seg_res`` [1, 9, 800, 500] seg refiner residual (before 8 tanh(x / 8)) |
| ``refiner.box_res`` [1, 8, 400, 250] box refiner residual (before 6 tanh(x / 6)) |
| ``refiner.e2e_ln / _res`` [1, 139] / [1, 36] e2e LayerNorm output / 3 tanh(mlp) |
| ========================== =========================== ==================================================== |
| |
| The 19 outputs are tapped under their ONNX names (``lane``, ``depth``, ``seg2d``, ``hm`` ... ``depth_mean``). |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any, Dict, Optional |
|
|
| import numpy as np |
|
|
| from ..ttaw.golden import NULL_TAPS, TapRegistry |
| from . import config as C |
| from .config import MeteorConfig, default_config |
| from .weights import MeteorWeights |
|
|
| __all__ = ["MeteorNet", "OUTPUT_NAMES", "child"] |
|
|
| OUTPUT_NAMES = ("lane", "depth", "seg2d", "hm", "reg", "hm2d_s0", "hm2d_s1", "hm2d_s2", "reg2d_s0", "reg2d_s1", |
| "reg2d_s2", "ego", "occ", "traj", "stationary", "tl", "risk", "lane_logit", "depth_mean") |
|
|
|
|
| def _torch(): |
| import torch |
|
|
| return torch |
|
|
|
|
| def child(module: str, index: int) -> str: |
| """ONNX scope of child ``index`` of a Sequential: ``child("seg_head/d1/d1.3", 0)`` = ``seg_head/d1/d1.3/d1.3.0``.""" |
| return f"{module}/{module.rsplit('/', 1)[-1]}.{index}" |
|
|
|
|
| class MeteorNet: |
| """The network as fp32 torch functions over the ONNX weights (no autograd, no nn.Module state).""" |
|
|
| def __init__(self, weights: MeteorWeights, cfg: Optional[MeteorConfig] = None, *, check_graph: bool = True): |
| torch = _torch() |
| self.cfg = (cfg or default_config()).check() |
| self.weights = weights |
| if check_graph: |
| weights.check_graph() |
| f32 = lambda a: torch.from_numpy(np.array(a, dtype=np.float32)) |
| self._conv: Dict[str, Any] = {} |
| for module, p in weights.iter_convs(): |
| pads = tuple(int(v) for v in p.pads) |
| if len(pads) != 4 or pads[:2] != pads[2:] or p.group != 1 or tuple(p.dilations) != (1, 1): |
| raise ValueError(f"{module}: unsupported conv attributes {p}") |
| self._conv[module] = (f32(p.weight), None if p.bias is None else f32(p.bias), |
| tuple(int(s) for s in p.strides), pads[:2]) |
| self._gemm: Dict[str, Any] = {} |
| for module in weights.gemm_modules(): |
| g = weights.gemm(module) |
| if g.trans_a or not g.trans_b or g.alpha != 1.0 or g.beta != 1.0: |
| raise ValueError(f"{module}: unsupported Gemm attributes") |
| self._gemm[module] = (f32(g.weight), f32(g.bias)) |
| self._pool = {name: f32(weights.pool_matrix(name)) for name in |
| ("ego_stem_w", "ego_stem_h", "tl_w", "tl_h", "fused_pool_w", "fused_pool_h", "sem_w", "sem_h", |
| "det_pool_w", "det_pool_h", "fused_mean_w", "fused_mean_h", "e2e_mean_w", "e2e_mean_h")} |
| self._mha = {name: weights.mha(name) for name in ("ego_attn", "agent_attn")} |
| self._mha_t = {name: (f32(m.in_weight), f32(m.in_bias), f32(m.out_weight), f32(m.out_bias)) |
| for name, m in self._mha.items()} |
| self._queries = {"ego": f32(weights.queries("ego")), "agent": f32(weights.queries("agent"))} |
| g, b, eps = weights.layer_norm() |
| self._ln = (f32(g), f32(b), float(eps)) |
| self._risk_gate = f32(np.asarray(weights.risk_gate(), np.float32).reshape(1)) |
| self._dec_gate = f32(np.asarray(weights.dec_gate(), np.float32).reshape(1)) |
| self._ground = f32(weights.lift_ground_points()) |
| zc = weights.zero_constants() |
| self._hist = f32(zc["hist"]) |
| self._mask = f32(zc["mask"]) |
| self._paint = torch.tensor(weights.paint_classes(), dtype=torch.long) |
| self._depth_centres = f32(C.depth_bin_centres(self.cfg.depth_mean_bins)).reshape(1, 1, C.DEPTH_BINS, 1, 1) |
|
|
| |
| def conv(self, x, module: str, relu: bool = False): |
| F = _torch().nn.functional |
| w, b, stride, pad = self._conv[module] |
| y = F.conv2d(x, w, b, stride=stride, padding=pad) |
| return F.relu(y) if relu else y |
|
|
| def block(self, x, module: str): |
| """ConvBlock = (3x3 conv, BN (folded), ReLU) x 2: children ``.0`` and ``.3`` (S:318).""" |
| return self.conv(self.conv(x, child(module, 0), True), child(module, 3), True) |
|
|
| def down(self, x, module: str): |
| """stride-2 3x3 conv + ReLU (child ``.0``) then a ConvBlock (child ``.3``): the encoder steps of the EDs.""" |
| return self.block(self.conv(x, child(module, 0), True), child(module, 3)) |
|
|
| def gemm(self, x, module: str): |
| w, b = self._gemm[module] |
| return _torch().addmm(b, x, w.t()) |
|
|
| @staticmethod |
| def resize(x, hw): |
| """ONNX Resize linear / half_pixel to a target size = F.interpolate(bilinear, align_corners=False).""" |
| F = _torch().nn.functional |
| return F.interpolate(x, size=tuple(int(v) for v in hw), mode="bilinear", align_corners=False) |
|
|
| def pool2(self, x, w_name: str, h_name: str): |
| """The exported adaptive pool / mean: (H_mat @ (x @ W_mat)) on the last two dims (S:356-366).""" |
| torch = _torch() |
| return torch.matmul(self._pool[h_name], torch.matmul(x, self._pool[w_name])) |
|
|
| @staticmethod |
| def nan_to_num(x): |
| """The graph's IsNaN / IsInf / Where chain (ONNX nodes 190-213 and 233-255): NaN -> 0, +inf -> 30, |
| -inf -> -30; the identity for finite values.""" |
| torch = _torch() |
| x = torch.where(torch.isnan(x), torch.zeros_like(x), x) |
| x = torch.where(torch.isinf(x) & (x > 0), torch.full_like(x, 30.0), x) |
| return torch.where(torch.isinf(x) & (x < 0), torch.full_like(x, -30.0), x) |
|
|
| def mha_pool(self, name: str, keys, queries): |
| """``nn.MultiheadAttention(q=learned queries, k=v=keys)`` as exported (torch 2.1, need_weights=False): |
| q and k|v projections, q / sqrt(head_dim), softmax over the 400 keys, out-projection. |
| ``keys`` [L, 1, E] (sequence first), ``queries`` [1, Q, E] -> [1, Q, E].""" |
| torch = _torch() |
| m = self._mha[name] |
| w_in, b_in, w_out, b_out = self._mha_t[name] |
| e, h = m.embed_dim, m.num_heads |
| hd = e // h |
| q_in = queries.transpose(0, 1) |
| q = b_in[:e] + torch.matmul(q_in, w_in[:e].t()) |
| kv = b_in[e:] + torch.matmul(keys, w_in[e:].t()) |
| k, v = kv[..., :e], kv[..., e:] |
| nq, nk = q.shape[0], k.shape[0] |
| qh = q.reshape(nq, h, hd).transpose(0, 1) |
| kh = k.reshape(nk, h, hd).transpose(0, 1) |
| vh = v.reshape(nk, h, hd).transpose(0, 1) |
| att = torch.softmax(torch.matmul(qh / m.scale_divisor, kh.transpose(1, 2)), dim=-1) |
| o = torch.matmul(att, vh).transpose(0, 1).reshape(nq, e) |
| o = torch.addmm(b_out, o, w_out.t()) |
| return o.reshape(nq, 1, e).transpose(0, 1) |
|
|
| |
| def normalize(self, imgs_u8, present=None): |
| """uint8 [1, 8, 3, 432, 768] -> [8, 3, 432, 768] float32: /255 (``onnx``) or (x/255 - mean)/std with absent |
| cameras zeroed AFTER normalisation (``imagenet``, SPEC section 2 imgs row).""" |
| torch = _torch() |
| x = imgs_u8.to(torch.float32) / 255.0 |
| if self.cfg.input_norm == "imagenet": |
| mean = torch.tensor(C.IMAGENET_MEAN, dtype=torch.float32).reshape(1, 1, 3, 1, 1) |
| std = torch.tensor(C.IMAGENET_STD, dtype=torch.float32).reshape(1, 1, 3, 1, 1) |
| x = (x - mean) / std |
| if present is not None: |
| keep = torch.as_tensor(np.asarray(present, np.float32)).reshape(1, -1, 1, 1, 1) |
| x = x * keep |
| return x.reshape(C.N_CAMS, 3, C.IMG_H, C.IMG_W) |
|
|
| def image_encoder(self, x, taps: TapRegistry = NULL_TAPS): |
| """ResNet-34 (torchvision layout, BN folded) + FPN -> f [8, 160, 108, 192] (S:301-302).""" |
| F = _torch().nn.functional |
| y = F.max_pool2d(self.conv(x, "stem/stem.0", True), kernel_size=3, stride=2, padding=1) |
| taps.tap("img.stem", y) |
| feats = [] |
| for li, nb in enumerate((3, 4, 6, 3), start=1): |
| for bi in range(nb): |
| p = f"layer{li}/layer{li}.{bi}" |
| ds = f"{p}/downsample/downsample.0" |
| idt = self.conv(y, ds) if ds in self._conv else y |
| y = F.relu(self.conv(self.conv(y, f"{p}/conv1", True), f"{p}/conv2") + idt) |
| feats.append(taps.tap(f"img.layer{li}", y)) |
| c1, c2, c3, c4 = feats |
| hw = (C.FEAT_H, C.FEAT_W) |
| s = self.conv(c1, "lat1") + self.resize(self.conv(c2, "lat2"), hw) |
| s = s + self.resize(self.conv(c3, "lat3"), hw) |
| s = s + self.resize(self.conv(c4, "lat4"), hw) |
| return taps.tap("img.fpn", self.conv(s, "fuse/fuse.0", True)) |
|
|
| def seg2d_head(self, f, taps: TapRegistry = NULL_TAPS): |
| """SegHeadED -> raw logits [8, 21, 108, 192] (S:303).""" |
| d1 = self.down(f, "seg_head/d1") |
| d2 = self.down(d1, "seg_head/d2") |
| m1 = self.block(d1 + self.resize(self.conv(d2, "seg_head/u1"), d1.shape[-2:]), "seg_head/m1") |
| u2 = self.resize(self.conv(m1, "seg_head/u2"), f.shape[-2:]) |
| y = self.conv(self.conv(self.conv(f, "seg_head/skip") + u2, "seg_head/out/out.0", True), "seg_head/out/out.3") |
| return taps.tap("seg2d.logits", y) |
|
|
| def depth_head(self, f, taps: TapRegistry = NULL_TAPS): |
| """4 ConvBlocks 160->128->128->96->64, 1x1 -> 64 bins, softmax over the bins (S:304).""" |
| torch = _torch() |
| y = f |
| for i in range(4): |
| y = self.block(y, f"depth_head/depth_head.{i}") |
| logits = taps.tap("depth.logits", self.conv(y, "depth_head/depth_head.4")) |
| return logits, taps.tap("depth.prob", torch.softmax(logits, dim=1)) |
|
|
| def det2d_heads(self, f): |
| """det2d stem + d8 / d16 pyramid, 1x1 heads at strides 4 / 8 / 16 (S:305) -> the six [1, 8, k, h, w] outputs.""" |
| s = self.block(f, "det2d_stem") |
| d8 = self.down(s, "det2d_d8") |
| d16 = self.down(d8, "det2d_d16") |
| out = {} |
| for i, (x, sfx) in enumerate(((s, ""), (d8, "8"), (d16, "16"))): |
| for kind, ch in (("hm2d", 10), ("reg2d", 4)): |
| y = self.conv(x, f"{kind}_head{sfx}") |
| out[f"{kind}_s{i}"] = y.reshape(1, C.N_CAMS, ch, y.shape[-2], y.shape[-1]) |
| return out |
|
|
| def tl_head(self, f): |
| """cat(f[cam 0], f[cam 6]) -> s2 conv + CB -> s2 conv + CB -> global mean (MatMul consts) -> Linear (S:306).""" |
| torch = _torch() |
| x = torch.cat([f[0:1], f[6:7]], dim=1) |
| x = self.block(self.conv(x, "tl_head/tl_head.0", True), "tl_head/tl_head.3") |
| x = self.block(self.conv(x, "tl_head/tl_head.4", True), "tl_head/tl_head.7") |
| x = self.pool2(x, "tl_w", "tl_h").flatten(1) |
| return self.gemm(x, "tl_fc") |
|
|
| def context(self, f, seg_logits, taps: TapRegistry = NULL_TAPS): |
| """ctx 1x1 + paint_proj(softmax(nan_to_num(raw seg logits))[paint classes]) (S:307; UNCLAMPED logits).""" |
| torch = _torch() |
| p = torch.softmax(self.nan_to_num(seg_logits), dim=1).index_select(1, self._paint) |
| return taps.tap("ctx.painted", self.conv(f, "ctx") + self.conv(p, "paint_proj")) |
|
|
| def lift(self, ctx, prob, K, T_cam_ego, taps: TapRegistry = NULL_TAPS): |
| """Depth-gated IPM lift to the 0.4 m grid, then bilinear x2 to the raw BEV (S:333-351, ONNX nodes 261-421).""" |
| torch = _torch() |
| F = torch.nn.functional |
| n = C.N_CAMS |
| T = T_cam_ego.reshape(n, 4, 4).to(torch.float32) |
| Km = K.reshape(n, 3, 3).to(torch.float32) |
| pc = torch.matmul(T, self._ground) |
| x, y, z = pc[:, 0], pc[:, 1], pc[:, 2] |
| zc = torch.clamp(z, min=C.LIFT_MIN_Z) |
| u = Km[:, 0, 0].unsqueeze(-1) * x / zc + Km[:, 0, 2].unsqueeze(-1) |
| v = Km[:, 1, 1].unsqueeze(-1) * y / zc + Km[:, 1, 2].unsqueeze(-1) |
| dist = torch.sqrt(x * x + y * y + z * z) |
| valid = ((z > C.LIFT_MIN_Z) & (u >= 0.0) & (u < float(C.IMG_W)) & (v >= 0.0) & (v < float(C.IMG_H)) |
| & (dist < C.LIFT_MAX_DIST)) |
| taps.tap("lift.valid", valid) |
| gu = torch.clamp(u / float(C.IMG_W - 1) * 2.0 - 1.0, -2.0, 2.0) |
| gv = torch.clamp(v / float(C.IMG_H - 1) * 2.0 - 1.0, -2.0, 2.0) |
| grid = taps.tap("lift.grid", torch.stack([gu, gv], dim=-1).unsqueeze(2)) |
| ctx_s = F.grid_sample(ctx, grid, mode="bilinear", padding_mode="zeros", align_corners=False).squeeze(3) |
| prob_s = F.grid_sample(prob, grid, mode="bilinear", padding_mode="zeros", align_corners=False).squeeze(3) |
| b = torch.clamp((dist - C.DEPTH_MIN) / C.DEPTH_STEP, 0.0, float(np.float32(C.LIFT_BIN_CLIP))) |
| b0 = torch.floor(b).to(torch.int64) |
| fr = (b - b0.to(torch.float32)).unsqueeze(1) |
| g0 = torch.gather(prob_s, 1, b0.unsqueeze(1)) |
| g1 = torch.gather(prob_s, 1, torch.clamp(b0 + 1, max=C.DEPTH_BINS - 1).unsqueeze(1)) |
| w = (g0 * (1.0 - fr) + g1 * fr + C.LIFT_W_EPS) * valid.unsqueeze(1).to(torch.float32) |
| taps.tap("lift.weight", w) |
| wn = w.reshape(1, n, 1, -1) |
| num = (ctx_s.reshape(1, n, C.CTX_C, -1) * wn).sum(dim=1) |
| den = torch.clamp(wn.sum(dim=1), min=C.LIFT_DEN_EPS) |
| bev04 = taps.tap("lift.bev", (num / den).reshape(1, C.CTX_C, C.LIFT_H, C.LIFT_W)) |
| return taps.tap("bev.raw", self.resize(bev04, (C.BEV_H, C.BEV_W))) |
|
|
| def temporal_fuse(self, raw, taps: TapRegistry = NULL_TAPS): |
| """No-history temporal fuse: fused = raw + tfuse(raw * 4 softmax(tgate(raw))[0]) (S:308).""" |
| torch = _torch() |
| gate = taps.tap("bev.gate", torch.softmax(self.conv(raw, "tgate_slim"), dim=1)[:, 0:1] * C.TGATE_SCALE) |
| t = self.block(self.conv(raw * gate, "tfuse3_slim/tfuse3_slim.0", True), "tfuse3_slim/tfuse3_slim.3") |
| return taps.tap("bev.fused", raw + t) |
|
|
| def lane_decoder(self, raw, taps: TapRegistry = NULL_TAPS): |
| """LaneDecED + lane_branch added to channels 4..6 -> lane_pre [1, 9, 800, 500] (S:309).""" |
| torch = _torch() |
| skip = self.block(raw, "dec/skip") |
| d1 = self.down(raw, "dec/d1") |
| d2 = self.down(d1, "dec/d2") |
| m1 = self.block(d1 + self.resize(self.conv(d2, "dec/u1"), d1.shape[-2:]), "dec/m1") |
| u2 = self.resize(self.conv(m1, "dec/u2"), skip.shape[-2:]) |
| out = self.conv(self.conv(skip + u2, "dec/out/out.0", True), "dec/out/out.3") |
| lb = self.conv(self.conv(self.conv(raw, "lane_branch/lane_branch.0", True), "lane_branch/lane_branch.3", True), |
| "lane_branch/lane_branch.6") |
| return taps.tap("lane.pre", torch.cat([out[:, :4], out[:, 4:7] + lb, out[:, 7:]], dim=1)) |
|
|
| def det_stem(self, raw, taps: TapRegistry = NULL_TAPS): |
| """_DetStemED on the raw BEV -> det_feat [1, 128, 400, 250], hm_pre, reg_pre (S:310).""" |
| s = self.down(raw, "det_stem/stem") |
| d8 = self.down(s, "det_stem/d8") |
| feat = self.conv(s + self.resize(self.conv(d8, "det_stem/u"), s.shape[-2:]), "det_stem/m/m.0", True) |
| taps.tap("det.feat", feat) |
| hm = taps.tap("det.hm_pre", self.conv(feat, "hm_head")) |
| return feat, hm, taps.tap("det.reg_pre", self.conv(feat, "reg_head")) |
|
|
| def occupancy(self, raw, taps: TapRegistry = NULL_TAPS): |
| """occ stem on the raw BEV crop rows 200:600, cols 50:450 -> 1x1 160 -> view [1, 10, 16, 200, 200] (S:311).""" |
| (r0, r1), (c0, c1) = C.OCC_CROP |
| x = self.down(raw[:, :, r0:r1, c0:c1], "occ_stem") |
| taps.tap("occ.feat", x) |
| y = self.conv(x, "occ_head") |
| return y.reshape(1, len(C.OCC_CLASSES), C.OCC_Z_BINS, y.shape[-2], y.shape[-1]) |
|
|
| def agent_forecast(self, raw, fused, det_feat, reg_pre, taps: TapRegistry = NULL_TAPS): |
| """traj_stem on cat(fused, mot = (raw - hist) * mask = 0), agent attention pool, traj head (S:312).""" |
| torch = _torch() |
| mot = (raw - self._hist) * self._mask |
| stem = self.block(self.conv(torch.cat([fused, mot], dim=1), "traj_stem/traj_stem.0", True), |
| "traj_stem/traj_stem.3") |
| del mot |
| taps.tap("traj.stem", stem) |
| reg46 = reg_pre[:, 4:6] |
| keys = self.pool2(det_feat, "det_pool_w", "det_pool_h").reshape(1, det_feat.shape[1], -1).permute(2, 0, 1) |
| att = self.mha_pool("agent_attn", keys, self._queries["agent"]) |
| delta = self.conv(att.mean(dim=1).reshape(1, -1, 1, 1), "agent_delta") |
| taps.tap("traj.agent_delta", delta) |
| feat = taps.tap("traj.feat", torch.cat([stem, det_feat], dim=1) + delta) |
| return self.conv(torch.cat([feat, reg46], dim=1), "traj_head") |
|
|
| def stationary(self, raw, taps: TapRegistry = NULL_TAPS): |
| """delta_stat on avgpool2(|raw - 0|) (S:313).""" |
| F = _torch().nn.functional |
| x = F.avg_pool2d((raw - self._hist).abs(), kernel_size=2, stride=2) |
| x = self.conv(self.conv(x, "delta_stat/delta_stat.1", True), "delta_stat/delta_stat.4", True) |
| taps.tap("stat.feat", x) |
| return self.conv(x, "delta_stat/delta_stat.7") |
|
|
| def risk(self, fused, taps: TapRegistry = NULL_TAPS): |
| """risk head on the fused BEV crop rows 200:600, cols 125:375 (S:314).""" |
| (r0, r1), (c0, c1) = C.RISK_CROP |
| x = self.block(self.block(fused[:, :, r0:r1, c0:c1], "risk_head/risk_head.0"), "risk_head/risk_head.1") |
| taps.tap("risk.feat", x) |
| return self.conv(x, "risk_head/risk_head.2") |
|
|
| def planner(self, fused, lane_pre, hm_pre, risk, v0, taps: TapRegistry = NULL_TAPS): |
| """Ego planner up to the e2e refiner (S:315): ego_stem + MLP, + ego_attn, + sem_ego, + kin, risk-aware |
| mode logits, decoupled decoder. Returns ``(ego_after_dec [1, 42], fused_mean [1, 96])``.""" |
| torch = _torch() |
| F = torch.nn.functional |
| with taps.scope("ego"): |
| e = self.block(self.conv(fused, "ego_stem/ego_stem.0", True), "ego_stem/ego_stem.3") |
| e = self.block(self.conv(e, "ego_stem/ego_stem.4", True), "ego_stem/ego_stem.7") |
| e = self.block(self.conv(e, "ego_stem/ego_stem.8", True), "ego_stem/ego_stem.11") |
| e = self.conv(e, "ego_stem/ego_stem.12", True) |
| e = taps.tap("stem", self.pool2(e, "ego_stem_w", "ego_stem_h").flatten(1)) |
| v0 = v0.reshape(1, 1).to(torch.float32) |
| h = torch.cat([e, v0], dim=1) |
| for i in (0, 2, 4): |
| h = F.relu(self.gemm(h, f"ego_mlp/ego_mlp.{i}")) |
| ego = taps.tap("mlp", self.gemm(h, "ego_mlp/ego_mlp.6")) |
| keys = self.pool2(fused, "fused_pool_w", "fused_pool_h").reshape(1, fused.shape[1], -1).permute(2, 0, 1) |
| att = self.mha_pool("ego_attn", keys, self._queries["ego"]) |
| ego = taps.tap("plus_attn", ego + taps.tap("attn", self.gemm(att.flatten(1), "ego_delta"))) |
| sem_in = torch.cat([F.avg_pool2d(torch.softmax(lane_pre, dim=1), kernel_size=2, stride=2), |
| torch.sigmoid(hm_pre)], dim=1) |
| taps.tap("sem_in", sem_in) |
| s = self.conv(self.conv(self.conv(sem_in, "sem_ego/sem_ego.0", True), "sem_ego/sem_ego.3", True), |
| "sem_ego/sem_ego.6", True) |
| s = self.pool2(s, "sem_w", "sem_h").flatten(1) |
| s = self.gemm(F.relu(self.gemm(s, "sem_ego/sem_ego.11")), "sem_ego/sem_ego.13") |
| ego = taps.tap("plus_sem", ego + taps.tap("sem", s)) |
| kin_in = torch.zeros(1, C.KIN_DIM, dtype=torch.float32) |
| kin_in[0, 0] = v0.reshape(()) / C.V0_SCALE |
| ego = taps.tap("plus_kin", ego + taps.tap("kin", self.gemm(kin_in, "kin_delta"))) |
| |
| paths = ego[:, :36].reshape(1, C.EGO_MODES, C.EGO_STEPS, 2) |
| gx = -paths[..., 1] / C.RISK_Y_HALF |
| gy = (C.RISK_X_HALF - paths[..., 0]) / C.RISK_X_HALF - 1.0 |
| samp = F.grid_sample(torch.sigmoid(risk), torch.stack([gx, gy], dim=-1), mode="bilinear", |
| padding_mode="border", align_corners=False) |
| samp = taps.tap("risk_sample", samp.mean(dim=3).squeeze(1)) |
| ego = torch.cat([ego[:, :36], ego[:, 36:39] - self._risk_gate * samp, ego[:, 39:]], dim=1) |
| taps.tap("after_risk", ego) |
| |
| fmean = self.pool2(fused, "fused_mean_w", "fused_mean_h").flatten(1) |
| dh = taps.tap("dec_head", self.gemm(fmean, "dec_head")).reshape(1, C.EGO_MODES, 12) |
| phi, vr = dh[..., :6], dh[..., 6:] |
| sp = torch.clamp(vr, min=0.0) + torch.log(1.0 + torch.exp(-vr.abs())) |
| step = torch.clamp(sp, max=C.DEC_SPEED_MAX) * 0.5 |
| wp = taps.tap("dec_wp", torch.stack([torch.cumsum(step * torch.cos(phi), dim=2), |
| torch.cumsum(step * torch.sin(phi), dim=2)], dim=-1)) |
| e36 = ego[:, :36].reshape(1, C.EGO_MODES, C.EGO_STEPS, 2) |
| e36 = e36 + torch.tanh(self._dec_gate) * (wp - e36) |
| ego = taps.tap("after_dec", torch.cat([e36.reshape(1, 36), ego[:, 36:]], dim=1)) |
| return ego, fmean |
|
|
| @staticmethod |
| def _range_channel(h: int, w: int): |
| """The refiners' constant range channel: row i -> (80 - (i + 0.5) * (160 / h)) / 80, float32 in graph order |
| (Range + 0.5, Reciprocal(h) * 160, Mul, 80 - ., / 80).""" |
| torch = _torch() |
| step = (1.0 / torch.tensor(float(h), dtype=torch.float32)) * 160.0 |
| rows = (80.0 - (torch.arange(h, dtype=torch.float32) + 0.5) * step) / 80.0 |
| return rows.reshape(1, 1, h, 1).expand(1, 1, h, w) |
|
|
| def seg_refiner(self, lane_pre, taps: TapRegistry = NULL_TAPS): |
| """BEVSegRefiner: lane_logit = clamp(lane_pre, +-20) + 8 tanh(res / 8) (S:316).""" |
| torch = _torch() |
| lc = torch.clamp(lane_pre, -C.SEG_REFINE_CLAMP, C.SEG_REFINE_CLAMP) |
| x = torch.cat([lc, self._range_channel(lc.shape[-2], lc.shape[-1])], dim=1) |
| p = "refiner/seg" |
| stem = self.conv(x, f"{p}/stem/stem.0", True) |
| d1 = self.down(stem, f"{p}/d1") |
| d2 = self.down(d1, f"{p}/d2") |
| d3 = self.down(d2, f"{p}/d3") |
| d4 = self.down(d3, f"{p}/d4") |
| m3 = self.block(d3 + self.resize(self.conv(d4, f"{p}/u4"), d3.shape[-2:]), f"{p}/m3") |
| m2 = self.block(d2 + self.resize(self.conv(m3, f"{p}/u3"), d2.shape[-2:]), f"{p}/m2") |
| m1 = self.block(d1 + self.resize(self.conv(m2, f"{p}/u2"), d1.shape[-2:]), f"{p}/m1") |
| o = stem + self.resize(self.conv(m1, f"{p}/u1"), stem.shape[-2:]) |
| res = taps.tap("refiner.seg_res", self.conv(self.conv(o, f"{p}/out/out.0", True), f"{p}/out/out.3")) |
| return lc + torch.tanh(res / C.SEG_REFINE_SCALE) * C.SEG_REFINE_SCALE |
|
|
| def box_refiner(self, hm_pre, reg_pre, taps: TapRegistry = NULL_TAPS): |
| """BEVBoxRefiner: hm, reg = clamp(.) + 6 tanh(res / 6) (S:317).""" |
| torch = _torch() |
| hc = torch.clamp(hm_pre, -C.BOX_HM_CLAMP, C.BOX_HM_CLAMP) |
| rc = torch.clamp(reg_pre, -C.BOX_REG_CLAMP, C.BOX_REG_CLAMP) |
| x = torch.cat([hc, rc, self._range_channel(hc.shape[-2], hc.shape[-1])], dim=1) |
| p = "refiner/box" |
| stem = self.conv(x, f"{p}/stem/stem.0", True) |
| d1 = self.down(stem, f"{p}/d1") |
| d2 = self.down(d1, f"{p}/d2") |
| d3 = self.down(d2, f"{p}/d3") |
| m2 = self.block(d2 + self.resize(self.conv(d3, f"{p}/u3"), d2.shape[-2:]), f"{p}/m2") |
| m1 = self.block(d1 + self.resize(self.conv(m2, f"{p}/u2"), d1.shape[-2:]), f"{p}/m1") |
| o = stem + self.resize(self.conv(m1, f"{p}/u1"), stem.shape[-2:]) |
| res = taps.tap("refiner.box_res", self.conv(self.conv(o, f"{p}/out/out.0", True), f"{p}/out/out.3")) |
| r = torch.tanh(res / C.BOX_REFINE_SCALE) * C.BOX_REFINE_SCALE |
| return hc + r[:, :2], rc + r[:, 2:] |
|
|
| def e2e_refiner(self, ego, v0, fused_mean, taps: TapRegistry = NULL_TAPS): |
| """E2ERefiner: ego[:36] += 3 tanh(mlp(LN([ego, v0, mean(fused)]))) (S:318; LN at the logical width 139).""" |
| torch = _torch() |
| F = torch.nn.functional |
| g, b, eps = self._ln |
| x = torch.cat([ego, v0.reshape(1, 1).to(torch.float32), fused_mean], dim=1) |
| h = taps.tap("refiner.e2e_ln", F.layer_norm(x, (x.shape[-1],), g, b, eps)) |
| h = F.relu(self.gemm(h, "refiner/e2e/mlp/mlp.0")) |
| h = F.relu(self.gemm(h, "refiner/e2e/mlp/mlp.2")) |
| r = taps.tap("refiner.e2e_res", torch.tanh(self.gemm(h, "refiner/e2e/mlp/mlp.4")) * C.E2E_SCALE) |
| return torch.cat([ego[:, :36] + r, ego[:, 36:]], dim=1) |
|
|
| def depth_mean(self, depth_logits): |
| """avgpool2 of the 64 logits -> softmax -> argmax -> +-2-bin window -> sum p c / max(sum p, 1e-6) -> fp16 |
| (ONNX nodes 1946-2014; bin centres per ``cfg.depth_mean_bins``).""" |
| torch = _torch() |
| F = torch.nn.functional |
| lp = F.avg_pool2d(depth_logits, kernel_size=2, stride=2) |
| p = torch.softmax(lp.reshape(1, C.N_CAMS, C.DEPTH_BINS, lp.shape[-2], lp.shape[-1]), dim=2) |
| am = torch.argmax(p, dim=2, keepdim=True) |
| r = torch.arange(C.DEPTH_BINS, dtype=torch.int64).reshape(1, 1, C.DEPTH_BINS, 1, 1) |
| pm = p * ((r - am).abs() <= 2).to(torch.float32) |
| num = (pm * self._depth_centres).sum(dim=2) |
| den = torch.clamp(pm.sum(dim=2), min=1e-6) |
| return (num / den).to(torch.float16) |
|
|
| |
| def forward(self, imgs, K, T_cam_ego, v0, *, present=None, taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]: |
| """The 19 outputs of the graph (torch tensors, ONNX dtypes) for one frame. |
| |
| ``imgs`` uint8 [1, 8, 3, 432, 768] RGB, ``K`` [1, 8, 3, 3] at 768x432, ``T_cam_ego`` [1, 8, 4, 4], ``v0`` [1]; |
| numpy or torch. ``present`` (8 flags) only matters for ``input_norm="imagenet"``.""" |
| torch = _torch() |
| t = lambda a, dt: a.to(dt) if isinstance(a, torch.Tensor) else torch.from_numpy(np.ascontiguousarray(a)).to(dt) |
| imgs = t(imgs, torch.uint8) |
| K = t(K, torch.float32) |
| T = t(T_cam_ego, torch.float32) |
| v0 = t(v0, torch.float32).reshape(1) |
| out: Dict[str, Any] = {} |
| with torch.inference_mode(): |
| x = taps.tap("img.input", self.normalize(imgs, present)) |
| f = self.image_encoder(x, taps) |
| del x |
| seg = self.seg2d_head(f, taps) |
| seg_c = taps.tap("seg2d.clamped", torch.clamp(self.nan_to_num(seg), -C.SEG2D_CLAMP, C.SEG2D_CLAMP)) |
| dlog, dprob = self.depth_head(f, taps) |
| ctx = self.context(f, seg, taps) |
| out.update(self.det2d_heads(f)) |
| out["tl"] = self.tl_head(f) |
| del f |
| raw = self.lift(ctx, dprob, K, T, taps) |
| del ctx, dprob |
| fused = self.temporal_fuse(raw, taps) |
| lane_pre = self.lane_decoder(raw, taps) |
| det_feat, hm_pre, reg_pre = self.det_stem(raw, taps) |
| out["occ"] = self.occupancy(raw, taps) |
| out["traj"] = self.agent_forecast(raw, fused, det_feat, reg_pre, taps) |
| out["stationary"] = self.stationary(raw, taps) |
| del raw, det_feat |
| out["risk"] = self.risk(fused, taps) |
| ego, fmean = self.planner(fused, lane_pre, hm_pre, out["risk"], v0, taps) |
| taps.tap("bev.fused_mean", fmean) |
| out["lane_logit"] = self.seg_refiner(lane_pre, taps) |
| out["hm"], out["reg"] = self.box_refiner(hm_pre, reg_pre, taps) |
| e2e_mean = self.pool2(fused, "e2e_mean_w", "e2e_mean_h").flatten(1) |
| out["ego"] = self.e2e_refiner(ego, v0, e2e_mean, taps) |
| del fused |
| out["lane"] = torch.argmax(out["lane_logit"], dim=1).to(torch.uint8) |
| dl5 = dlog.reshape(1, C.N_CAMS, C.DEPTH_BINS, C.FEAT_H, C.FEAT_W) |
| out["depth"] = torch.argmax(dl5, dim=2).to(torch.uint8) |
| out["seg2d"] = torch.argmax(seg_c.reshape(1, C.N_CAMS, len(C.SEG2D_CLASSES), C.FEAT_H, C.FEAT_W), |
| dim=2).to(torch.uint8) |
| out["depth_mean"] = self.depth_mean(dlog) |
| for name in OUTPUT_NAMES: |
| taps.tap(name, out[name]) |
| return {name: out[name] for name in OUTPUT_NAMES} |
|
|
| __call__ = forward |
|
|
