File size: 23,831 Bytes
fe51a63 | 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 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 | # SPDX-License-Identifier: Apache-2.0
"""The deployed ``frnet_<variant>.onnx`` of AutowareFoundation/lidar_frnet@v2.0 read as DATA, the way the TT port
reads it: the vendored ``ttaw.weights.OnnxWeights`` (initializers addressed by the node that consumes them; nothing
is executed). 120 of the 186 initializer *names* differ between the OT128 and QT128 exports (``onnx::Conv_1487`` vs
``onnx::Conv_1485``, S:frnet:98-101), node names do not, so every weight is found from its node.
Where BatchNorm is folded (float64, one final rounding to float32; ``ttaw.weights.fold_bn_linear``):
- every point layer ``MatMul (no bias) -> BatchNormalization(eps 1e-3) -> Relu`` (``ffe_layers.0-2``,
``point_stem``, ``point_fusion_layers.0-3``, ``point_fuse_layer1-2``, ``mlps.0-3``) becomes one
``y = relu(x @ W' + b')`` -- exactly the port's ``ttnn.linear(..., bias, activation="relu")``;
- ``pre_norm`` (BatchNorm1d over the 8 decorated features) is kept as ``(scale, shift)`` float32 (``x * scale +
shift``): the port applies it on the host with the feature decoration (PLAN.md 2.8), so the first device layer is
a plain linear;
- the 53 convs already carry the exporter's BN fold (Conv with bias, S:frnet:245); the three Gemms
(``ffe_layers.3``, ``compression_layers.0``, ``conv_seg``) have their own bias.
The structure is recovered from the graph, never re-typed: stage count, BasicBlocks per stage, strides, the
frustum grid of every frustum2pixel (the ``ConstantOfShape`` shape of its ``ScatterND``), every activation (the
conv's consumer), and every wiring assumption of the reference model (concat orders, skip connections, the
decoder's resize inputs, the head's residual skips, the Softmax axis) is asserted on the graph
(:func:`load_weights` raises ``ValueError`` when one does not hold). Each layer records the ONNX tensor that holds
its output, so ``reference/ort.py`` exposes the same taps from ONNX Runtime.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple
import numpy as np
from ..ttaw.weights import BatchNorm, OnnxWeights, file_sha256, fold_bn_linear
from .config import BLOCKS, NUM_CLASSES, ONNX_SHA256, STRIDES, variant_files
__all__ = ["Linear", "Conv", "Stage", "FRNetWeights", "load_weights", "check_variant_grid", "MODULE_TAPS",
"OPT_IN_TAPS"]
VE = "/model/voxel_encoder"
BB = "/model/backbone"
@dataclass(frozen=True)
class Linear:
"""``y = act(x @ weight + bias)``; ``weight`` (in, out) float32 (BN folded where the graph has one)."""
name: str
weight: np.ndarray
bias: np.ndarray
relu: bool
onnx_node: str
onnx_output: str # ONNX tensor equal to this layer's output (after the ReLU when there is one)
bn_node: Optional[str] = None
@property
def in_features(self) -> int:
return int(self.weight.shape[0])
@property
def out_features(self) -> int:
return int(self.weight.shape[1])
@dataclass(frozen=True)
class Conv:
"""``y = act(conv2d(x, weight, bias, stride, padding))`` (NCHW, BN already folded by the exporter)."""
name: str
weight: np.ndarray # (Cout, Cin, k, k) float32
bias: np.ndarray # (Cout,) float32
stride: int
padding: int
activation: Optional[str] # "hardswish" | "sigmoid" | None (a residual add follows)
onnx_node: str
onnx_output: str # ONNX tensor equal to the output (after the activation when there is one)
@property
def kernel(self) -> int:
return int(self.weight.shape[-1])
@property
def in_channels(self) -> int:
return int(self.weight.shape[1])
@property
def out_channels(self) -> int:
return int(self.weight.shape[0])
@dataclass(frozen=True)
class Stage:
"""One backbone stage ``layer<index>`` (1-based) plus its point / pixel fusion (``*_layers.<index - 1>``)."""
index: int
stride: int # cumulative stride of the stage output (1, 2, 4, 8)
blocks: Tuple[Tuple[Conv, Conv, Optional[Conv]], ...] # (conv1 + HardSwish, conv2, downsample) per BasicBlock
block_outputs: Tuple[str, ...] # ONNX tensor of each block output (after add + HardSwish)
point_fusion: Linear # [pixel2point(x_s), p] 256 -> 128 + BN + ReLU
pixel_fusion: Conv # [frustum2pixel(segmax(p)), x_s] 256 -> 128 + HardSwish
attention: Tuple[Conv, Conv] # 128 -> 128 + HardSwish, 128 -> 128 + Sigmoid
segmax_output: str # ONNX tensor: per-voxel max of point_fusion's output (M, 128)
output: str # ONNX tensor: x = pixel_fusion * attention + x_s
@dataclass(frozen=True)
class FRNetWeights:
variant: str
frustum_hw: Tuple[int, int]
strides: Tuple[int, ...]
pre_norm: BatchNorm
pre_norm_scale: np.ndarray # (8,) float32: pre_norm(x) == x * scale + shift
pre_norm_shift: np.ndarray
ffe: Tuple[Linear, ...] # 8 -> 64 -> 128 -> 256 (ReLU, BN folded), 256 -> 256 (Gemm, no activation)
compression: Linear # 256 -> 16 + ReLU (on the per-voxel max)
stem: Tuple[Conv, ...] # 16 -> 64 -> 128 -> 128, each + HardSwish
point_stem: Linear # [pixel2point(stem), pf3] 384 -> 128 + BN + ReLU
point_stem_segmax: str # ONNX tensor of the per-voxel max of point_stem (M, 128)
fusion_stem: Conv # [frustum2pixel(segmax(p0)), stem] 256 -> 128 + HardSwish
stages: Tuple[Stage, ...]
fuse_layer1: Conv # concat of the 5 up-sampled maps 640 -> 256 + HardSwish
fuse_layer2: Conv # 256 -> 128 + HardSwish
point_fuse: Tuple[Linear, Linear] # 640 -> 256, 256 -> 128 (+ BN + ReLU)
head_mlps: Tuple[Linear, ...] # 128 -> 128, 128 -> 256, 256 -> 128, 128 -> 64 (+ BN + ReLU), then + skip
conv_seg: Linear # 64 -> 27
onnx_taps: Dict[str, str] # tap name -> ONNX tensor (module and opt-in taps)
onnx_path: str = ""
onnx_sha256: str = ""
raw_parameters: int = 0 # learned values in the ONNX initializers
folded_parameters: int = field(default=0) # values after the BN folds (what the device holds)
@property
def convs(self) -> List[Conv]:
out = list(self.stem) + [self.fusion_stem]
for st in self.stages:
for c1, c2, ds in st.blocks:
out += [c1, c2] + ([ds] if ds is not None else [])
out += [st.pixel_fusion, *st.attention]
return out + [self.fuse_layer1, self.fuse_layer2]
@property
def linears(self) -> List[Linear]:
return [*self.ffe, self.compression, self.point_stem, *(s.point_fusion for s in self.stages),
*self.point_fuse, *self.head_mlps, self.conv_seg]
@property
def num_classes(self) -> int:
return self.conv_seg.out_features
# Module taps of the reference (tests/test_reference_cpu.py compares every one with ORT; the TT port's per-module
# gates use the same names). 1-based stage numbers: stage l = ONNX layer<l> + *_layers.<l - 1>.
MODULE_TAPS: Tuple[str, ...] = (
"encoder.input", "encoder.pf0", "encoder.pf1", "encoder.pf2", "encoder.pf3", "encoder.voxel_max",
"encoder.voxel_feats", "backbone.frustum", "backbone.stem", "backbone.point_stem", "backbone.fusion_stem",
*(f"backbone.{kind}{i}" for i in range(1, 5)
for kind in ("layer", "point_fusion", "pixel_fusion", "attention", "stage")),
"decoder.fuse1", "decoder.fuse2", "decoder.point_fuse1", "decoder.point_fuse2",
*(f"head.mlp{i}" for i in range(4)), "head.logits", "pred_probs",
)
# Opt-in taps (TapRegistry include patterns): every BasicBlock output, the per-voxel max of every point branch
# (the K1 / segmented-max sites), the five up-sampled decoder inputs.
OPT_IN_TAPS: Tuple[str, ...] = (
"backbone.layer*.block*", "backbone.segmax*", "decoder.up*",
)
def _fail(msg: str) -> None:
raise ValueError(f"frnet ONNX structure: {msg}")
def _single(w: OnnxWeights, tensor: str, op_type: str) -> Any:
try:
return w.consumer_of(tensor, op_type)
except KeyError as e:
_fail(str(e))
def _point_layer(w: OnnxWeights, matmul: str, name: str) -> Linear:
mm = w.node(matmul)
if mm.op_type != "MatMul" or not w.has(mm.inputs[1]) or w.has(mm.inputs[0]):
_fail(f"{matmul} is not x @ W with a constant W")
bn = _single(w, mm.outputs[0], "BatchNormalization")
relu = _single(w, bn.outputs[0], "Relu")
wf, bf = fold_bn_linear(w.matmul_weight(matmul), None, w.batchnorm(bn.name), layout="in_out")
return Linear(name, wf, bf, True, mm.name, relu.outputs[0], bn.name)
def _gemm(w: OnnxWeights, node: str, name: str) -> Linear:
g = w.gemm(node)
if g.trans_a or g.alpha != 1.0 or g.beta != 1.0 or g.bias is None:
_fail(f"{node}: unexpected Gemm attributes")
weight = np.asarray(g.weight, np.float32)
weight = weight.T if g.trans_b else weight
out = w.node(node).outputs[0]
relus = [c for c in w.consumers(out) if c.op_type == "Relu"]
relu = len(relus) == 1 and len(w.consumers(out)) == 1
return Linear(name, np.ascontiguousarray(weight), np.asarray(g.bias, np.float32), relu, node,
relus[0].outputs[0] if relu else out)
def _conv(w: OnnxWeights, node: str, name: str) -> Conv:
c = w.conv(node)
k = c.kernel_shape
if (c.group != 1 or set(c.dilations) != {1} or len(k) != 2 or k[0] != k[1] or len(set(c.strides)) != 1
or len(set(c.pads)) != 1 or c.bias is None or c.auto_pad not in ("NOTSET", "")):
_fail(f"{node}: unexpected Conv attributes {c.kernel_shape} {c.strides} {c.pads} {c.dilations} {c.group}")
out = w.node(node).outputs[0]
cons = w.consumers(out)
act, tensor = None, out
if len(cons) == 1 and cons[0].op_type in ("HardSwish", "Sigmoid"):
act, tensor = cons[0].op_type.lower(), cons[0].outputs[0]
return Conv(name, np.asarray(c.weight, np.float32), np.asarray(c.bias, np.float32), int(c.strides[0]),
int(c.pads[0]), act, node, tensor)
def _producer_op(w: OnnxWeights, tensor: str) -> str:
p = w.producer(tensor)
return p.op_type if p is not None else ""
def _scatter_max_of(w: OnnxWeights, tensor: str) -> str:
found = [c for c in w.consumers(tensor) if c.op_type == "ScatterElements" and c.attrs.get("reduction") == "max"]
if len(found) != 1:
_fail(f"{tensor} does not feed exactly one ScatterElements(max)")
return found[0].outputs[0]
def load_weights(model_dir: Any, variant: str, *, check_sha256: bool = True) -> FRNetWeights:
"""Parse ``frnet_<variant>.onnx`` of a weights directory (module docstring)."""
path = variant_files(model_dir, variant)["onnx"]
digest = file_sha256(path)
if check_sha256 and digest != ONNX_SHA256[variant]: # refused before it is parsed
raise ValueError(f"{path.name}: sha256 {digest} is not the pinned v2.0 file {ONNX_SHA256[variant]}")
w = OnnxWeights(path)
if w.input_names != ("points", "coors", "voxel_coors", "inverse_map") or w.output_names != ("pred_probs",):
_fail(f"inputs {w.input_names} / outputs {w.output_names}")
taps: Dict[str, str] = {}
# ---- frustum grids of the six frustum2pixel (ScatterND) sites, in graph order
grids = []
for sc in w.nodes("ScatterND"):
shape_node = w.producer(w.producer(sc.inputs[0]).inputs[0])
dims = [w.array(t).tolist() for t in shape_node.inputs[:3] if w.has(t)]
if len(dims) != 3 or dims[0] != [1]:
_fail(f"{sc.name}: cannot read its (1, H, W, C) image shape")
grids.append((int(dims[1][0]), int(dims[2][0])))
if len(grids) != 2 + len(STRIDES):
_fail(f"{len(grids)} ScatterND sites, expected {2 + len(STRIDES)}")
H, W = grids[0]
strides = tuple(H // h for h, _ in grids[2:])
if grids[1] != (H, W) or strides != STRIDES or any(W // s != g[1] or H // s != g[0]
for s, g in zip(strides, grids[2:])):
_fail(f"frustum grids {grids} do not follow strides {STRIDES}")
# ---- encoder
bn_pre = w.find_node(f"{VE}/pre_norm/*", "BatchNormalization")
pre = w.batchnorm(bn_pre.name)
taps["encoder.input"] = bn_pre.outputs[0]
ffe = [_point_layer(w, f"{VE}/ffe_layers.{i}/ffe_layers.{i}.0/MatMul", f"encoder.ffe{i}") for i in range(3)]
for i in range(1, 3):
if w.node(ffe[i].onnx_node).inputs[0] != ffe[i - 1].onnx_output:
_fail(f"ffe_layers.{i} does not consume ffe_layers.{i - 1}")
if w.node(ffe[0].onnx_node).inputs[0] != bn_pre.outputs[0]:
_fail("ffe_layers.0 does not consume pre_norm")
ffe3 = _gemm(w, f"{VE}/ffe_layers.3/Gemm", "encoder.ffe3")
if ffe3.relu or w.node(ffe3.onnx_node).inputs[0] != ffe[2].onnx_output:
_fail("ffe_layers.3 is not a plain Gemm on ffe_layers.2")
ffe.append(ffe3)
for i, lin in enumerate(ffe):
taps[f"encoder.pf{i}"] = lin.onnx_output
taps["encoder.voxel_max"] = _scatter_max_of(w, ffe3.onnx_output)
compression = _gemm(w, f"{VE}/compression_layers/compression_layers.0/Gemm", "encoder.compression")
if not compression.relu or w.node(compression.onnx_node).inputs[0] != taps["encoder.voxel_max"]:
_fail("compression_layers is not Gemm + ReLU on the per-voxel max")
taps["encoder.voxel_feats"] = compression.onnx_output
# ---- stem and point stem
stem = tuple(_conv(w, f"{BB}/stem/stem.{i}/Conv", f"backbone.stem{j}") for j, i in enumerate((0, 3, 6)))
if any(c.activation != "hardswish" or c.stride != 1 for c in stem):
_fail("stem convs must be stride 1 + HardSwish")
taps["backbone.frustum"] = w.node(stem[0].onnx_node).inputs[0]
if _producer_op(w, taps["backbone.frustum"]) != "Transpose":
_fail("the stem does not read the transposed frustum image")
taps["backbone.stem"] = stem[-1].onnx_output
point_stem = _point_layer(w, f"{BB}/point_stem/point_stem.0/MatMul", "backbone.point_stem")
cat = w.producer(w.node(point_stem.onnx_node).inputs[0])
if cat.op_type != "Concat" or len(cat.inputs) != 2 or cat.inputs[1] != ffe3.onnx_output:
_fail("point_stem input is not concat[pixel2point(stem), ffe_layers.3]")
taps["backbone.point_stem"] = point_stem.onnx_output
point_stem_segmax = _scatter_max_of(w, point_stem.onnx_output)
taps["backbone.segmax0"] = point_stem_segmax
fusion_stem = _conv(w, f"{BB}/fusion_stem/fusion_stem.0/Conv", "backbone.fusion_stem")
cat = w.producer(w.node(fusion_stem.onnx_node).inputs[0])
if fusion_stem.activation != "hardswish" or cat.op_type != "Concat" or cat.inputs[1] != stem[-1].onnx_output:
_fail("fusion_stem is not HardSwish(conv(concat[frustum2pixel(p0), stem]))")
taps["backbone.fusion_stem"] = fusion_stem.onnx_output
# ---- stages
stages: List[Stage] = []
x_in, p_prev = fusion_stem.onnx_output, point_stem.onnx_output
for li in range(1, 5):
conv1s = w.nodes("Conv", f"{BB}/layer{li}/layer{li}.*/conv1/Conv")
if len(conv1s) != BLOCKS[li - 1]:
_fail(f"layer{li} has {len(conv1s)} BasicBlocks, expected {BLOCKS[li - 1]}")
blocks, outs = [], []
for bi in range(len(conv1s)):
pre_b = f"{BB}/layer{li}/layer{li}.{bi}"
c1 = _conv(w, f"{pre_b}/conv1/Conv", f"backbone.layer{li}.block{bi}.conv1")
c2 = _conv(w, f"{pre_b}/conv2/Conv", f"backbone.layer{li}.block{bi}.conv2")
ds = None
if w.nodes("Conv", f"{pre_b}/downsample/*"):
ds = _conv(w, f"{pre_b}/downsample/downsample.0/Conv", f"backbone.layer{li}.block{bi}.downsample")
if c1.activation != "hardswish" or c2.activation is not None or (ds is not None and ds.activation):
_fail(f"{pre_b}: unexpected activations")
if w.node(c1.onnx_node).inputs[0] != x_in or (ds is not None and w.node(ds.onnx_node).inputs[0] != x_in):
_fail(f"{pre_b} does not read the previous block output")
add = _single(w, c2.onnx_output, "Add")
idt = ds.onnx_output if ds is not None else x_in
if set(add.inputs) != {c2.onnx_output, idt}:
_fail(f"{pre_b}: the residual add does not take conv2 + identity")
hs = _single(w, add.outputs[0], "HardSwish")
blocks.append((c1, c2, ds))
outs.append(hs.outputs[0])
taps[f"backbone.layer{li}.block{bi}"] = hs.outputs[0]
x_in = hs.outputs[0]
x_s = outs[-1]
taps[f"backbone.layer{li}"] = x_s
pf = _point_layer(w, f"{BB}/point_fusion_layers.{li - 1}/point_fusion_layers.{li - 1}.0/MatMul",
f"backbone.point_fusion{li}")
cat = w.producer(w.node(pf.onnx_node).inputs[0])
if cat.op_type != "Concat" or len(cat.inputs) != 2 or cat.inputs[1] != p_prev:
_fail(f"point_fusion_layers.{li - 1} input is not concat[pixel2point(x), p]")
seg = _scatter_max_of(w, pf.onnx_output)
taps[f"backbone.point_fusion{li}"] = pf.onnx_output
taps[f"backbone.segmax{li}"] = seg
pix = _conv(w, f"{BB}/pixel_fusion_layers.{li - 1}/pixel_fusion_layers.{li - 1}.0/Conv",
f"backbone.pixel_fusion{li}")
cat = w.producer(w.node(pix.onnx_node).inputs[0])
if pix.activation != "hardswish" or cat.op_type != "Concat" or cat.inputs[1] != x_s:
_fail(f"pixel_fusion_layers.{li - 1} is not HardSwish(conv(concat[frustum2pixel(p), x]))")
a0 = _conv(w, f"{BB}/attention_layers.{li - 1}/attention_layers.{li - 1}.0/Conv",
f"backbone.attention{li}.conv0")
a1 = _conv(w, f"{BB}/attention_layers.{li - 1}/attention_layers.{li - 1}.3/Conv",
f"backbone.attention{li}.conv1")
if (a0.activation != "hardswish" or a1.activation != "sigmoid" or w.node(a0.onnx_node).inputs[0] !=
pix.onnx_output or w.node(a1.onnx_node).inputs[0] != a0.onnx_output):
_fail(f"attention_layers.{li - 1} is not sigmoid(conv(hswish(conv(pixel_fusion))))")
mul = _single(w, a1.onnx_output, "Mul")
if set(mul.inputs) != {pix.onnx_output, a1.onnx_output}:
_fail(f"stage {li}: the attention product is not pixel_fusion * attention")
add = _single(w, mul.outputs[0], "Add")
if set(add.inputs) != {mul.outputs[0], x_s}:
_fail(f"stage {li}: the output is not pixel_fusion * attention + x")
taps[f"backbone.pixel_fusion{li}"] = pix.onnx_output
taps[f"backbone.attention{li}"] = a1.onnx_output
taps[f"backbone.stage{li}"] = add.outputs[0]
stride = STRIDES[li - 1]
stages.append(Stage(li, stride, tuple(blocks), tuple(outs), pf, pix, (a0, a1), seg, add.outputs[0]))
x_in, p_prev = add.outputs[0], pf.onnx_output
got = [1] + [st.blocks[0][0].stride for st in stages[1:]]
if tuple(int(np.prod(got[: i + 1])) for i in range(4)) != STRIDES:
_fail(f"stage strides {got} do not give {STRIDES}")
# ---- decoder
fuse1 = _conv(w, f"{BB}/fuse_layer1/conv/Conv", "decoder.fuse1")
fuse2 = _conv(w, f"{BB}/fuse_layer2/conv/Conv", "decoder.fuse2")
cat = w.producer(w.node(fuse1.onnx_node).inputs[0])
sources = [fusion_stem.onnx_output] + [st.output for st in stages]
if cat.op_type != "Concat" or len(cat.inputs) != 5:
_fail("fuse_layer1 does not read a 5-way concat")
for k, (t, src) in enumerate(zip(cat.inputs, sources)):
rs = w.producer(t)
if rs.op_type != "Resize" or rs.inputs[0] != src or rs.attrs.get("mode") != "linear" or \
rs.attrs.get("coordinate_transformation_mode") != "align_corners":
_fail(f"decoder input {k} is not an align_corners bilinear Resize of {src}")
taps[f"decoder.up{k}"] = t
if fuse1.activation != "hardswish" or fuse2.activation != "hardswish" or \
w.node(fuse2.onnx_node).inputs[0] != fuse1.onnx_output:
_fail("fuse_layer1 / fuse_layer2 are not HardSwish convs in sequence")
taps["decoder.fuse1"] = fuse1.onnx_output
taps["decoder.fuse2"] = fuse2.onnx_output
pf1 = _point_layer(w, f"{BB}/point_fuse_layer1/point_fuse_layer1.0/MatMul", "decoder.point_fuse1")
pf2 = _point_layer(w, f"{BB}/point_fuse_layer2/point_fuse_layer2.0/MatMul", "decoder.point_fuse2")
cat = w.producer(w.node(pf1.onnx_node).inputs[0])
if cat.op_type != "Concat" or list(cat.inputs) != [point_stem.onnx_output] + [s.point_fusion.onnx_output
for s in stages]:
_fail("point_fuse_layer1 input is not concat[p0, p1, p2, p3, p4]")
if w.node(pf2.onnx_node).inputs[0] != pf1.onnx_output:
_fail("point_fuse_layer2 does not consume point_fuse_layer1")
taps["decoder.point_fuse1"] = pf1.onnx_output
taps["decoder.point_fuse2"] = pf2.onnx_output
# ---- head: h = pixel2point(fuse2); h = mlp_i(h) + skip_i; logits; softmax
mlps = [_point_layer(w, f"/model/mlps.{i}/mlps.{i}.0/MatMul", f"head.mlp{i}") for i in range(4)]
skips = [pf2.onnx_output, ffe[2].onnx_output, ffe[1].onnx_output, ffe[0].onnx_output]
prev = None
for i, (lin, skip) in enumerate(zip(mlps, skips)):
add = _single(w, lin.onnx_output, "Add")
if set(add.inputs) != {lin.onnx_output, skip}:
_fail(f"head mlps.{i} skip is not {skip}")
if prev is not None and w.node(lin.onnx_node).inputs[0] != prev:
_fail(f"head mlps.{i} does not consume the previous residual sum")
taps[f"head.mlp{i}"] = add.outputs[0]
prev = add.outputs[0]
conv_seg = _gemm(w, "/model/conv_seg/Gemm", "head.logits")
if conv_seg.relu or w.node(conv_seg.onnx_node).inputs[0] != prev or conv_seg.out_features != NUM_CLASSES:
_fail("conv_seg is not the final 64 -> 27 Gemm on the last residual sum")
sm = _single(w, conv_seg.onnx_output, "Softmax")
if sm.attrs.get("axis", -1) not in (-1, 1) or sm.outputs[0] != "pred_probs":
_fail("the output is not Softmax(axis=-1) of conv_seg")
taps["head.logits"] = conv_seg.onnx_output
taps["pred_probs"] = "pred_probs"
scale = pre.scale.astype(np.float32)
shift = pre.shift.astype(np.float32)
raw = int(sum(int(np.asarray(w.array(t)).size) for t in w.initializer_names()
if np.issubdtype(np.asarray(w.array(t)).dtype, np.floating)))
out = FRNetWeights(variant=variant, frustum_hw=(H, W), strides=strides, pre_norm=pre, pre_norm_scale=scale,
pre_norm_shift=shift, ffe=tuple(ffe), compression=compression, stem=stem,
point_stem=point_stem, point_stem_segmax=point_stem_segmax, fusion_stem=fusion_stem,
stages=tuple(stages), fuse_layer1=fuse1, fuse_layer2=fuse2, point_fuse=(pf1, pf2),
head_mlps=tuple(mlps), conv_seg=conv_seg, onnx_taps=taps, onnx_path=str(path),
onnx_sha256=digest, raw_parameters=raw)
folded = sum(c.weight.size + c.bias.size for c in out.convs) + sum(
lin.weight.size + lin.bias.size for lin in out.linears) + 16
object.__setattr__(out, "folded_parameters", int(folded))
missing = [t for t in MODULE_TAPS if t not in taps]
if missing:
_fail(f"no ONNX tensor for taps {missing}")
return out
def check_variant_grid(weights: FRNetWeights, frustum_h: int, frustum_w: int) -> None:
"""The ml_package grid must equal the grid baked into the ONNX (S:frnet:95-98: 1024 OT128 / 256 QT128)."""
if weights.frustum_hw != (frustum_h, frustum_w):
raise ValueError(f"ml_package frustum {frustum_h}x{frustum_w} != the ONNX's {weights.frustum_hw}")
|