code/tt_sceneseg/reference/model.py

# 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()}