File size: 10,126 Bytes
aec3f16 | 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 | # 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()}
|