| |
| """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"] |
|
|
| |
| 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 |
| 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)} |
|
|
