# SPDX-License-Identifier: Apache-2.0 """ONNX Runtime oracle of the deployed graph: the 19 outputs plus any internal tensor as a named tap. Only for tests and golden generation in the research venv (onnxruntime is not a runtime dependency of the bundle, BUNDLE_CONVENTIONS.md section 13). The graph is modified in memory only (extra graph outputs; for ``input_norm="imagenet"`` the training normalisation inserted after ``/Div``, the oracle of SPEC section 10 risk 1, re-implemented from ``research/meteor/scripts/norm_ab_test.py``); the downloaded file is never written. ORT on this graph is bit-reproducible at a fixed thread count (SPEC section 6). """ from __future__ import annotations import os from typing import Any, Dict, Iterable, Mapping, Optional, Sequence import numpy as np from . import config as C __all__ = ["ORT_TAPS", "OrtOracle", "available"] # reference tap name -> ONNX tensor name (shapes as in the graph; see reference/model.py) ORT_TAPS: Dict[str, str] = { "img.input": "/net/Reshape_output_0", "img.stem": "/net/stem/stem.3/MaxPool_output_0", "img.layer1": "/net/layer1/layer1.2/relu_1/Relu_output_0", "img.layer2": "/net/layer2/layer2.3/relu_1/Relu_output_0", "img.layer3": "/net/layer3/layer3.5/relu_1/Relu_output_0", "img.layer4": "/net/layer4/layer4.2/relu_1/Relu_output_0", "img.fpn": "/net/fuse/fuse.2/Relu_output_0", "seg2d.logits": "/net/seg_head/out/out.3/Conv_output_0", "seg2d.clamped": "/net/Clip_output_0", "depth.logits": "/net/depth_head/depth_head.4/Conv_output_0", "depth.prob": "/net/Softmax_output_0", "ctx.painted": "/net/Add_3_output_0", "lift.grid": "/net/Unsqueeze_6_output_0", "lift.valid": "/net/And_8_output_0", "lift.weight": "/net/Mul_10_output_0", "lift.bev": "/net/Reshape_5_output_0", "bev.raw": "/net/Resize_3_output_0", "bev.gate": "/net/Mul_12_output_0", "bev.fused": "/net/Add_11_output_0", "lane.pre": "/net/Concat_7_output_0", "det.feat": "/net/det_stem/m/m.2/Relu_output_0", "det.hm_pre": "/net/hm_head/Conv_output_0", "det.reg_pre": "/net/reg_head/Conv_output_0", "occ.feat": "/net/occ_stem/occ_stem.3/occ_stem.3.5/Relu_output_0", "traj.stem": "/net/traj_stem/traj_stem.3/traj_stem.3.5/Relu_output_0", "traj.agent_delta": "/net/agent_delta/Conv_output_0", "traj.feat": "/net/Add_15_output_0", "stat.feat": "/net/delta_stat/delta_stat.6/Relu_output_0", "risk.feat": "/net/risk_head/risk_head.1/risk_head.1.5/Relu_output_0", "ego.stem": "/net/Flatten_output_0", "ego.mlp": "/net/ego_mlp/ego_mlp.6/Gemm_output_0", "ego.attn": "/net/ego_delta/Gemm_output_0", "ego.plus_attn": "/net/Add_13_output_0", "ego.sem_in": "/net/Concat_23_output_0", "ego.sem": "/net/sem_ego/sem_ego.13/Gemm_output_0", "ego.plus_sem": "/net/Add_14_output_0", "ego.kin": "/net/kin_delta/Gemm_output_0", "ego.plus_kin": "/net/Add_17_output_0", "ego.risk_sample": "/net/Squeeze_2_output_0", "ego.after_risk": "/net/ScatterND_1_output_0", "ego.dec_head": "/net/dec_head/Gemm_output_0", "ego.dec_wp": "/net/Concat_36_output_0", "ego.after_dec": "/net/ScatterND_2_output_0", "bev.fused_mean": "/net/Flatten_3_output_0", "refiner.seg_res": "/net/refiner/seg/out/out.3/Conv_output_0", "refiner.box_res": "/net/refiner/box/out/out.3/Conv_output_0", "refiner.e2e_ln": "/net/refiner/e2e/norm/LayerNormalization_output_0", "refiner.e2e_res": "/net/refiner/e2e/Mul_output_0", } def available() -> bool: try: import onnxruntime # noqa: F401 (probe only) except ImportError: return False return True def _insert_imagenet(model) -> None: """(x / 255 - mean) / std right after ``/Div`` (in memory; SPEC section 10 risk 1).""" from onnx import helper, numpy_helper g = model.graph div = next(n for n in g.node if n.name == "/Div") out = div.output[0] mean = np.array(C.IMAGENET_MEAN, np.float32).reshape(1, 1, 3, 1, 1) std = np.array(C.IMAGENET_STD, np.float32).reshape(1, 1, 3, 1, 1) g.initializer.extend([numpy_helper.from_array(mean, "tt_meteor_mean"), numpy_helper.from_array(std, "tt_meteor_std")]) for n in g.node: if n is not div: for i, x in enumerate(n.input): if x == out: n.input[i] = "tt_meteor_norm" idx = list(g.node).index(div) g.node.insert(idx + 1, helper.make_node("Sub", [out, "tt_meteor_mean"], ["tt_meteor_sub"], name="/tt_meteor/Sub")) g.node.insert(idx + 2, helper.make_node("Div", ["tt_meteor_sub", "tt_meteor_std"], ["tt_meteor_norm"], name="/tt_meteor/Div")) class OrtOracle: """``OrtOracle(onnx_path, taps=[...], threads=4)``; ``oracle(feed) -> {name: array}`` with the 19 outputs under their ONNX names and the requested taps under their reference names (``ORT_TAPS``; raw ONNX tensor names are accepted too).""" def __init__(self, onnx_path: os.PathLike, taps: Iterable[str] = (), *, threads: int = 4, input_norm: str = "onnx"): import onnx import onnxruntime as ort self.taps = {t: ORT_TAPS.get(t, t) for t in taps} model = onnx.load(str(onnx_path)) if input_norm == "imagenet": _insert_imagenet(model) elif input_norm != "onnx": raise ValueError(f"input_norm {input_norm!r}") have = {o.name for o in model.graph.output} for name in self.taps.values(): if name not in have: model.graph.output.append(onnx.helper.make_empty_tensor_value_info(name)) so = ort.SessionOptions() so.intra_op_num_threads = int(threads) so.inter_op_num_threads = 1 self.session = ort.InferenceSession(model.SerializeToString(), sess_options=so, providers=["CPUExecutionProvider"]) del model self.output_names = [o.name for o in self.session.get_outputs()] self.version = ort.__version__ self.threads = int(threads) def __call__(self, feed: Mapping[str, np.ndarray], *, outputs: Optional[Sequence[str]] = None) -> Dict[str, Any]: inv = {v: k for k, v in self.taps.items()} feed = {k: np.asarray(feed[k]) for k in ("imgs", "K", "T_cam_ego", "v0")} res = self.session.run(outputs, feed) names = list(outputs) if outputs is not None else self.output_names return {inv.get(n, n): v for n, v in zip(names, res)}