File size: 13,072 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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 | # SPDX-License-Identifier: Apache-2.0
"""Pure-PyTorch fp32 re-implementation of the plugin-free parts of ``bevdet_one_lt_d.onnx`` with per-module taps.
Ported from the verified research reference (``research/bevdet/scripts/bevdet_torch.py``: PCC 1.0 against ONNX
Runtime on the deployed sub-graphs, SPEC section 6). Module boundaries mirror the training code (LCH1238/BEVDet@export:
ResNet-50, CustomFPN, DepthNet (BEVDepth), CustomResNet, FPN_LSS, CenterHead) so that each TT stage has a reference
tap. Every convolution runs with the weights and the attributes of its ONNX node (``reference.weights``).
``BEVDetNet.image_encoder(img, cam_params)``: normalised images [6, 3, 256, 704] + ``cam_params`` [1, 6, 27] ->
(depth [6, 118, 16, 44] after the softmax over depth bins, feat [6, 80, 16, 44]); ``gates=`` replaces the camera MLP /
SE branch by precomputed per-camera gates (the TT port's form, ``reference.port_form.camera_gates``).
``BEVDetNet.bev_decoder(bev_in)``: GatherBEV output [1, 720, 128, 128] -> the 6 head maps [1, k, 128, 128].
Tap names (``ttaw.golden.TapRegistry``; the device PCC gates use the same names):
============================= ===================== ============================================================
tap shape where
============================= ===================== ============================================================
``img.backbone.stem`` [6, 64, 64, 176] conv1 7x7 s2 + ReLU + MaxPool 3x3 s2
``img.backbone.layer1..4`` [6, 256..2048, ...] ResNet-50 stages (layer3 = C4 [6,1024,16,44], layer4 = C5)
``img.neck`` [6, 512, 16, 44] CustomFPN output (fpn conv, no activation)
``depthnet.reduce`` [6, 512, 16, 44] reduce_conv 3x3 + ReLU
``depthnet.gate_ctx/_depth`` [6, 512] sigmoid(SE(MLP(BN1d(cam_params)))), calibration-only
``feat`` [6, 80, 16, 44] context_conv(reduce * gate_ctx): the BEVPool features
``depthnet.blocks.0..2`` [6, 512, 16, 44] BasicBlocks on reduce * gate_depth
``depthnet.aspp.b1..b4/.gap`` [6, 96, 16, 44] ASPP branches (dilation 1/6/12/18, GAP: [6, 96, 1, 1])
``depthnet.aspp`` [6, 512, 16, 44] ASPP conv1 (480 -> 512) + ReLU
``depthnet.logits`` [6, 118, 16, 44] depth_conv.4 1x1
``depth`` [6, 118, 16, 44] softmax over the 118 bins
``bev.backbone.0..2`` [1, 160, 64, 64] ... CustomResNet stages (F0, F1, F2 [1, 640, 16, 16])
``bev.neck.up4`` [1, 640, 64, 64] bilinear x4 (align_corners) of F2
``bev.neck.conv`` [1, 512, 64, 64] two 3x3 + ReLU on cat(F0, up4)
``bev.neck.up2`` [1, 512, 128, 128] bilinear x2 (align_corners)
``bev.neck`` [1, 256, 128, 128] 3x3 + ReLU, 1x1
``head.shared`` [1, 64, 128, 128] shared 3x3 + ReLU
``head.<name>.hidden`` [1, 64, 128, 128] per-branch 3x3 + ReLU
``head.reg_0`` ... ``heatmap_0`` [1, k, 128, 128] the ONNX outputs (``HEAD_NAMES``)
============================= ===================== ============================================================
"""
from __future__ import annotations
from typing import Any, Dict, Optional, Tuple
import numpy as np
from ..ttaw.golden import NULL_TAPS, TapRegistry
from .config import HEAD_NAMES, BEVDetConfig, default_config
from .weights import HEAD_BRANCHES, BEVDetWeights
__all__ = ["BEVDetNet"]
def _torch():
import torch
return torch
class BEVDetNet:
"""The network as fp32 torch functions over the ONNX weights (no autograd, no nn.Module state)."""
def __init__(self, weights: BEVDetWeights, cfg: Optional[BEVDetConfig] = None, *, dtype: Any = None):
torch = _torch()
self.cfg = cfg or default_config()
self.cfg.check()
self.weights = weights
self.dtype = dtype or torch.float32
self._params: Dict[str, Tuple[Any, Optional[Any], Any]] = {}
for module, conv in weights.iter_convs():
w = torch.from_numpy(np.array(conv.weight, np.float32)).to(self.dtype) # own copy (writable)
b = None if conv.bias is None else torch.from_numpy(np.array(conv.bias, np.float32)).to(self.dtype)
self._params[module] = (w, b, conv)
dn = "img_view_transformer.depth_net."
self._lin = {m: tuple(torch.tensor(np.array(a, np.float32), dtype=self.dtype) for a in weights.linear(dn + m))
for m in ("context_mlp.fc1", "context_mlp.fc2", "depth_mlp.fc1", "depth_mlp.fc2")}
bn = weights.depth_bn()
self._bn = tuple(torch.from_numpy(np.asarray(a, np.float32)).to(self.dtype)
for a in (bn.gamma, bn.beta, bn.mean, bn.var))
self._bn_eps = float(np.float32(bn.eps))
# ---- primitives -----------------------------------------------------------------------------------
def conv(self, x: Any, module: str, relu: bool = False) -> Any:
torch = _torch()
import torch.nn.functional as F
w, b, c = self._params[module]
y = F.conv2d(x, w, b, stride=c.stride, padding=c.padding, dilation=c.dilation, groups=c.groups)
return torch.relu(y) if relu else y
def linear(self, x: Any, module: str, relu: bool = False) -> Any:
torch = _torch()
w, b = self._lin[module]
y = torch.nn.functional.linear(x, w, b)
return torch.relu(y) if relu else y
# ---- image encoder --------------------------------------------------------------------------------
def backbone(self, x: Any, taps: TapRegistry = NULL_TAPS) -> Tuple[Any, Any]:
"""ResNet-50 (pytorch style: stride on the 3x3, 1x1 downsample, BN folded) -> (C4, C5)."""
torch = _torch()
import torch.nn.functional as F
x = self.conv(x, "img_backbone.conv1", relu=True)
x = taps.tap("img.backbone.stem", F.max_pool2d(x, 3, 2, 1))
outs = []
for li, nb in enumerate([3, 4, 6, 3], start=1):
for bi in range(nb):
p = f"img_backbone.layer{li}.{bi}"
idt = self.conv(x, p + ".downsample.0") if bi == 0 else x
y = self.conv(x, p + ".conv1", relu=True)
y = self.conv(y, p + ".conv2", relu=True)
y = self.conv(y, p + ".conv3")
x = torch.relu(y + idt)
outs.append(taps.tap(f"img.backbone.layer{li}", x))
return outs[2], outs[3]
def neck(self, c4: Any, c5: Any, taps: TapRegistry = NULL_TAPS) -> Any:
"""CustomFPN: L0 = 1x1(C4) + nearest-upsample(1x1(C5)) (ONNX Resize nearest / asymmetric / floor), then the
3x3 fpn conv with no activation."""
import torch.nn.functional as F
l0 = self.conv(c4, "img_neck.lateral_convs.0.conv")
l1 = self.conv(c5, "img_neck.lateral_convs.1.conv")
l0 = l0 + F.interpolate(l1, size=l0.shape[2:], mode="nearest")
return taps.tap("img.neck", self.conv(l0, "img_neck.fpn_convs.0.conv"))
def camera_mlp_input(self, cam_params: Any) -> Any:
"""``depth_net.bn`` (BatchNorm1d(27), float32 like the graph) on ``cam_params.reshape(-1, 27)``."""
torch = _torch()
g, b, m, v = self._bn
x = cam_params.reshape(-1, cam_params.shape[-1]).to(self.dtype)
return (x - m) / torch.sqrt(v + self._bn_eps) * g + b
def camera_gates(self, cam_params: Any) -> Tuple[Any, Any]:
"""The deployed graph's camera branch: ``sigmoid(SE(MLP(BN1d(cam_params))))`` for context and depth ->
([6, 512], [6, 512]) float32."""
torch = _torch()
dn = "img_view_transformer.depth_net."
mlp_in = self.camera_mlp_input(cam_params)
def gate(mlp: str, se: str) -> Any:
v = self.linear(self.linear(mlp_in, mlp + ".fc1", relu=True), mlp + ".fc2")[..., None, None]
g = self.conv(self.conv(v, dn + se + ".conv_reduce", relu=True), dn + se + ".conv_expand")
return torch.sigmoid(g)[..., 0, 0]
return gate("context_mlp", "context_se"), gate("depth_mlp", "depth_se")
def depth_net(self, x: Any, cam_params: Any = None, *, gates: Optional[Tuple[Any, Any]] = None,
taps: TapRegistry = NULL_TAPS) -> Tuple[Any, Any]:
"""DepthNet (BEVDepth, camera-aware): -> (depth softmax [6, 118, 16, 44], feat [6, 80, 16, 44])."""
torch = _torch()
import torch.nn.functional as F
dn = "img_view_transformer.depth_net."
x = taps.tap("depthnet.reduce", self.conv(x, dn + "reduce_conv.0", relu=True))
if gates is None:
gates = self.camera_gates(cam_params)
g_ctx = taps.tap("depthnet.gate_ctx", gates[0])
g_depth = taps.tap("depthnet.gate_depth", gates[1])
feat = taps.tap("feat", self.conv(x * g_ctx[:, :, None, None], dn + "context_conv"))
d = x * g_depth[:, :, None, None]
for bi in range(3):
q = dn + f"depth_conv.{bi}"
y = self.conv(d, q + ".conv1", relu=True)
y = self.conv(y, q + ".conv2")
d = taps.tap(f"depthnet.blocks.{bi}", torch.relu(y + d))
q = dn + "depth_conv.3"
x1 = taps.tap("depthnet.aspp.b1", self.conv(d, q + ".aspp1.atrous_conv", relu=True))
x2 = taps.tap("depthnet.aspp.b2", self.conv(d, q + ".aspp2.atrous_conv", relu=True))
x3 = taps.tap("depthnet.aspp.b3", self.conv(d, q + ".aspp3.atrous_conv", relu=True))
x4 = taps.tap("depthnet.aspp.b4", self.conv(d, q + ".aspp4.atrous_conv", relu=True))
x5 = self.conv(F.adaptive_avg_pool2d(d, 1), q + ".global_avg_pool.1", relu=True)
taps.tap("depthnet.aspp.gap", x5)
x5 = F.interpolate(x5, size=x4.shape[2:], mode="bilinear", align_corners=True)
d = taps.tap("depthnet.aspp", self.conv(torch.cat([x1, x2, x3, x4, x5], 1), q + ".conv1", relu=True))
logits = taps.tap("depthnet.logits", self.conv(d, dn + "depth_conv.4"))
return taps.tap("depth", logits.softmax(dim=1)), feat
def image_encoder(self, img: Any, cam_params: Any = None, *, gates: Optional[Tuple[Any, Any]] = None,
taps: TapRegistry = NULL_TAPS) -> Tuple[Any, Any]:
"""[6, 3, 256, 704] normalised images (+ ``cam_params`` [1, 6, 27] or precomputed ``gates``) ->
(depth [6, 118, 16, 44], feat [6, 80, 16, 44])."""
torch = _torch()
with torch.no_grad():
img = torch.as_tensor(img).to(self.dtype)
cp = None if cam_params is None else torch.as_tensor(cam_params).to(self.dtype)
if gates is not None:
gates = tuple(torch.as_tensor(g).to(self.dtype) for g in gates)
c4, c5 = self.backbone(img, taps)
x = self.neck(c4, c5, taps)
return self.depth_net(x, cp, gates=gates, taps=taps)
# ---- BEV decoder ----------------------------------------------------------------------------------
def bev_encoder(self, x: Any, taps: TapRegistry = NULL_TAPS) -> Any:
"""CustomResNet (BasicBlocks, 3x3 s2 downsample) + FPN_LSS (bilinear align_corners x4 / x2)."""
torch = _torch()
import torch.nn.functional as F
feats = []
for li in range(3):
for bi in range(2):
p = f"img_bev_encoder_backbone.layers.{li}.{bi}"
idt = self.conv(x, p + ".downsample") if bi == 0 else x
y = self.conv(x, p + ".conv1", relu=True)
y = self.conv(y, p + ".conv2")
x = torch.relu(y + idt)
feats.append(taps.tap(f"bev.backbone.{li}", x))
x2 = taps.tap("bev.neck.up4", F.interpolate(feats[2], scale_factor=4, mode="bilinear", align_corners=True))
x = torch.cat([feats[0], x2], 1)
x = self.conv(x, "img_bev_encoder_neck.conv.0", relu=True)
x = taps.tap("bev.neck.conv", self.conv(x, "img_bev_encoder_neck.conv.3", relu=True))
x = taps.tap("bev.neck.up2", F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True))
x = self.conv(x, "img_bev_encoder_neck.up2.1", relu=True)
return taps.tap("bev.neck", self.conv(x, "img_bev_encoder_neck.up2.4"))
def center_head(self, x: Any, taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]:
x = taps.tap("head.shared", self.conv(x, "pts_bbox_head.shared_conv.conv", relu=True))
out = {}
for h, name in zip(HEAD_BRANCHES, HEAD_NAMES):
p = f"pts_bbox_head.task_heads.0.{h}"
y = taps.tap(f"head.{name}.hidden", self.conv(x, p + ".0.conv", relu=True))
out[name] = taps.tap(f"head.{name}", self.conv(y, p + ".1"))
return out
def bev_decoder(self, bev_in: Any, taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]:
"""[1, 720, 128, 128] -> {``reg_0``: [1, 2, 128, 128], ..., ``heatmap_0``: [1, 10, 128, 128]}."""
torch = _torch()
with torch.no_grad():
x = torch.as_tensor(bev_in).to(self.dtype)
return self.center_head(self.bev_encoder(x, taps), taps)
|