# SPDX-License-Identifier: Apache-2.0 """fp32 PyTorch reference of the deployed SceneSeg graph, module by module, with per-module taps. The graph is the one VisionPilot runs (``SceneSeg_FP32.onnx``; S:sceneseg 4.2), written out as functions over the layer weights of :mod:`.weights` (loaded from the ONNX by consuming node, BN already folded), so the TT graph can mirror it module by module; ``tests/test_reference_cpu.py`` proves it against onnxruntime on the shipped ONNX per tap. - **Backbone**: torchvision EfficientNet-B0 ``features`` (stem 3x3 s2 + SiLU; 16 MBConv blocks: [expand 1x1 + SiLU], depthwise kxk + SiLU, SqueezeExcitation (mean, fc1 + SiLU, fc2 + Sigmoid, scale), project 1x1, residual when stride 1 and cin == cout; head 1x1 + SiLU). Taps ``l0`` .. ``l8`` (the stage outputs; SceneSeg uses l0, l2, l3, l4, l8). - **SceneContext**: ``v = mean_HW(l8)``; Linear 1280->800 + GELU, 800->800 + GELU, 800->200 + Sigmoid; reshape to 1x1x10x20; Conv3x3 1->128->256->512->1280, each + GELU (``gate``); ``context = gate * l8 + l8``. - **SceneNeck / SceneSegHead**: five ``ConvTranspose2d(k2, s2)`` up-samplings, each + a 1x1 skip conv of a backbone tap (l4, l3, l2, l0; none after the fifth), then two Conv3x3 + GELU; the final Conv3x3 64 -> 3 gives the logits. GELU is the exact erf form (the ONNX exports ``x * (erf(x / sqrt 2) + 1) * 0.5``); SiLU is ``x * sigmoid(x)``. Usage:: ref = SceneSegReference.from_onnx(onnx_path) logits = ref.run_numpy(x_nchw_f32)["logits"] # [1, 3, 320, 640] float32 taps = TapRegistry(); ref.run_numpy(x, taps=taps) # + every tap of config.TAPS (numpy float32) """ from __future__ import annotations from pathlib import Path from typing import Any, Dict, Union import numpy as np from . import config from .weights import ARCH, LayerWeights, load_state_dict_weights, load_weights __all__ = ["SceneSegReference"] class SceneSegReference: """The SceneSeg network in fp32 torch (batch 1, 320x640). Thread count: whatever ``torch.set_num_threads`` says (keep <= 4 on the shared host).""" def __init__(self, weights: Dict[str, LayerWeights]): import torch missing = [s.key for s in ARCH if s.key not in weights] if missing: raise ValueError(f"weights for {len(missing)} layers missing, e.g. {missing[:3]}") self.weights = weights self._w = {k: (torch.from_numpy(lw.weight.copy()), torch.from_numpy(lw.bias.copy())) for k, lw in weights.items()} self._spec = {k: lw.spec for k, lw in weights.items()} @classmethod def from_onnx(cls, onnx_path: Union[str, Path], *, check_sha256: bool = True) -> "SceneSegReference": return cls(load_weights(onnx_path, check_sha256=check_sha256)) @classmethod def from_state_dict(cls, path: Union[str, Path]) -> "SceneSegReference": """From the upstream ``.pth`` / the re-hosted ``.safetensors`` (BN folded in fp64).""" return cls(load_state_dict_weights(path)) # ---- primitives ---------------------------------------------------------------------------------------- @staticmethod def _act(y, act: str): import torch import torch.nn.functional as F if act == "silu": return F.silu(y) if act == "gelu": return F.gelu(y) # exact erf GELU (approximate="none") if act == "sigmoid": return torch.sigmoid(y) return y def layer(self, key: str, x): """One weighted layer of :data:`ARCH` followed by its activation.""" import torch.nn.functional as F w, b = self._w[key] s = self._spec[key] if s.kind == "conv": y = F.conv2d(x, w, b, stride=s.stride, padding=s.pad, groups=s.groups) elif s.kind == "conv_transpose": y = F.conv_transpose2d(x, w, b, stride=s.stride) else: y = F.linear(x, w, b) return self._act(y, s.act) @staticmethod def _tap(taps, name: str, value): if taps is not None: taps.tap(name, value) return value # ---- modules ------------------------------------------------------------------------------------------- def mbconv(self, stage: int, block: int, x): """One MBConv block (torchvision ``efficientnet.MBConv``) of backbone stage ``stage``.""" e, k, stride, cin, cout, _ = config.B0_SETTINGS[stage - 1] cin_b = cin if block == 0 else cout key = f"Backbone.encoder.{stage}.{block}.block" i = 0 y = x if e != 1: y = self.layer(f"{key}.0.0", y) i = 1 y = self.layer(f"{key}.{i}.0", y) # depthwise kxk + SiLU s = self.layer(f"{key}.{i + 1}.fc2", self.layer(f"{key}.{i + 1}.fc1", y.mean(dim=(2, 3), keepdim=True))) y = self.layer(f"{key}.{i + 2}.0", y * s) # SE scale, then project 1x1 if (stride if block == 0 else 1) == 1 and cin_b == cout: y = y + x return y def backbone(self, image, taps=None) -> Dict[str, Any]: """``image`` [1, 3, 320, 640] -> the stage outputs ``l0`` .. ``l8``.""" feats = {"l0": self._tap(taps, "l0", self.layer("Backbone.encoder.0.0", image))} y = feats["l0"] for s, (_, _, _, _, _, n) in enumerate(config.B0_SETTINGS, start=1): for b in range(n): y = self.mbconv(s, b, y) feats[f"l{s}"] = self._tap(taps, f"l{s}", y) feats["l8"] = self._tap(taps, "l8", self.layer("Backbone.encoder.8.0", y)) return feats def context(self, l8, taps=None): """SceneContext: the global MLP prior that gates the stride-32 features (``scene_context.py:14-56``).""" v = self._tap(taps, "context.mean", l8.mean(dim=(2, 3))) # [1, 1280] c = self.layer("SceneContext.context_layer_0", v) c = self.layer("SceneContext.context_layer_1", c) c = self._tap(taps, "context.mlp", self.layer("SceneContext.context_layer_2", c)) # [1, 200], sigmoid g = c.reshape(1, 1, *config.CONTEXT_HW) # batch 1, 10 x 20 (as exported) for i in range(3, 7): g = self.layer(f"SceneContext.context_layer_{i}", g) g = self._tap(taps, "context.gate", g) return self._tap(taps, "context", g * l8 + l8) def _decoder_block(self, mod: str, block, x, skip, taps, names): up, _, skip_mod, _, decodes = block y = self.layer(f"{mod}.{up}", x) if skip_mod is not None: y = y + self.layer(f"{mod}.{skip_mod}", skip) y = self._tap(taps, names[0], y) for name, _, _ in decodes: y = self.layer(f"{mod}.{name}", y) return self._tap(taps, names[1], y) def neck(self, context, feats, taps=None): """SceneNeck (``scene_neck.py``): three up-sampling blocks with the l4, l3, l2 skips -> 256 x 80 x 160.""" names = (("neck.up0", "neck.dec1"), ("neck.up1", "neck.dec3"), ("neck.up2", "neck")) y = context for block, skip, nm in zip(config.NECK_BLOCKS, ("l4", "l3", "l2"), names): y = self._decoder_block("SceneNeck", block, y, feats[skip], taps, nm) return y def head(self, neck, feats, taps=None): """SceneSegHead (``scene_seg_head.py``): up-sampling with the l0 skip, a plain up-sampling, final 3x3.""" y = self._decoder_block("SceneSegHead", config.HEAD_BLOCKS[0], neck, feats["l0"], taps, ("head.up3", "head.dec7")) up, _, _, _, decodes = config.HEAD_BLOCKS[1] y = self._tap(taps, "head.up4", self.layer(f"SceneSegHead.{up}", y)) y = self._tap(taps, "head.dec8", self.layer(f"SceneSegHead.{decodes[0][0]}", y)) y = self._tap(taps, "head.dec9", self.layer(f"SceneSegHead.{decodes[1][0]}", y)) return self._tap(taps, "logits", self.layer(f"SceneSegHead.{config.FINAL_CONV[0]}", y)) def forward(self, image, taps=None): """``image`` float32 [1, 3, 320, 640] (the ONNX input) -> logits [1, 3, 320, 640].""" if tuple(image.shape) != (1, 3) + tuple(config.INPUT_HW): raise ValueError(f"input must be [1, 3, {config.INPUT_HW[0]}, {config.INPUT_HW[1]}] (batch 1, fixed " f"size: the SceneContext reshape), got {tuple(image.shape)}") feats = self.backbone(image, taps) context = self.context(feats["l8"], taps) neck = self.neck(context, feats, taps) return self.head(neck, feats, taps) __call__ = forward def run_numpy(self, x: np.ndarray, taps=None) -> Dict[str, np.ndarray]: """numpy float32 NCHW in -> ``{"logits": numpy float32}`` (taps, if given, are filled as numpy).""" import torch with torch.inference_mode(): logits = self.forward(torch.from_numpy(np.ascontiguousarray(x, np.float32)), taps) return {"logits": logits.numpy()} def run_from(self, start: str, inputs: Dict[str, np.ndarray], taps=None) -> Dict[str, np.ndarray]: """Teacher-forced partial runs (TT module gates): start at a module boundary with the given reference taps. ``start``: ``"context"`` (needs ``l8``), ``"neck"`` (``context``, ``l2``, ``l3``, ``l4``), ``"head"`` (``neck``, ``l0``). Returns ``{"logits": ...}``; ``taps`` collects the downstream taps.""" import torch t = {k: torch.from_numpy(np.ascontiguousarray(v, np.float32)) for k, v in inputs.items()} with torch.inference_mode(): if start == "context": ctx = self.context(t["l8"], taps) t["context"] = ctx start = "neck" if start == "neck": t["neck"] = self.neck(t["context"], t, taps) start = "head" if start != "head": raise ValueError(f"start must be 'context', 'neck' or 'head', got {start!r}") logits = self.head(t["neck"], t, taps) return {"logits": logits.numpy()}