# SPDX-License-Identifier: Apache-2.0 """ONNX Runtime on the shipped v5.0 ONNX files: the oracle of the CPU reference (research and tests only; onnxruntime is not a runtime dependency of the bundle or the image). ``OrtPlanner.run(raw)`` is the Autoware ``multi_step`` pipeline with ORT as the backend (a port of the verified research script ``research/diffusion-planner/scripts/dp_reference.py``): the node's normalization and speed masks, encoder.onnx, the DPM-Solver++(2M) loop of :mod:`tt_diffusion_planner.host.solver` calling decoder.onnx 11 times, then turn_indicator.onnx. ``taps=[...]`` exposes intermediate tensors of the encoder / decoder graphs by ONNX tensor name (the decoder taps are recorded per evaluation), for the per-module PCC test of the reference. """ from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, List, Sequence import numpy as np from . import config as C from ..host.normalize import normalize_inputs, speed_masks from ..host.solver import apply_prefix_constraint, dpm_solver_sample from ..host.features import decoder_masks from ..ttaw import io as tio from .weights import load_param_json __all__ = ["OrtPlanner", "OrtResult", "available"] def available() -> bool: import importlib.util return importlib.util.find_spec("onnxruntime") is not None def _session(model: Any, threads: int, optimize: bool = True): import onnxruntime as ort so = ort.SessionOptions() so.graph_optimization_level = (ort.GraphOptimizationLevel.ORT_ENABLE_ALL if optimize else ort.GraphOptimizationLevel.ORT_DISABLE_ALL) so.intra_op_num_threads = threads so.inter_op_num_threads = 1 so.log_severity_level = 3 src = model.SerializeToString() if hasattr(model, "SerializeToString") else str(model) return ort.InferenceSession(src, so, providers=["CPUExecutionProvider"]) def _with_outputs(path: Path, names: Sequence[str]): """The model with extra graph outputs (intermediate tensors by name).""" import onnx m = onnx.load(str(path)) have = {o.name for o in m.graph.output} m.graph.output.extend([onnx.ValueInfoProto(name=n) for n in names if n not in have]) return m @dataclass class OrtResult: norm: Dict[str, np.ndarray] encoding: np.ndarray # [1, 564, 256] final_x0: np.ndarray # [1, 321, 81, 4] logit: np.ndarray # [1, 5] denoising_steps: List[np.ndarray] denoising_timesteps: List[float] eval_times: List[float] eval_inputs: List[np.ndarray] = field(default_factory=list) eval_outputs: List[np.ndarray] = field(default_factory=list) taps: Dict[str, np.ndarray] = field(default_factory=dict) class OrtPlanner: """``OrtPlanner(weights_dir, threads=4, encoder_taps=(...), decoder_taps=(...))``. With taps the sessions run without graph optimizations (every intermediate kept as computed); without, ``ORT_ENABLE_ALL`` like ``onnxruntime_inference.cpp:172``.""" def __init__(self, weights_dir: Path, *, threads: int = 4, encoder_taps: Sequence[str] = (), decoder_taps: Sequence[str] = (), turn_taps: Sequence[str] = ()): wd = Path(weights_dir) self.normalization = load_param_json(wd / C.PARAM_JSON) self.encoder_taps, self.decoder_taps, self.turn_taps = list(encoder_taps), list(decoder_taps), list(turn_taps) opt = not (encoder_taps or decoder_taps or turn_taps) self.enc = _session(_with_outputs(wd / C.ENCODER_ONNX, self.encoder_taps), threads, opt) self.dec = _session(_with_outputs(wd / C.DECODER_ONNX, self.decoder_taps), threads, opt) self.turn = _session(_with_outputs(wd / C.TURN_INDICATOR_ONNX, self.turn_taps), threads, opt) def normalize(self, raw: Any) -> Dict[str, np.ndarray]: arrays = tio.load_named_arrays(raw, C.INPUT_SCHEMA) norm = normalize_inputs(arrays, self.normalization.observation) norm.update(speed_masks(norm)) return norm def encode(self, norm: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: feed = {k: norm[k] for k in C.ENCODER_INPUTS} outs = self.enc.run(["encoding"] + self.encoder_taps, feed) return dict(zip(["encoding"] + self.encoder_taps, outs)) def decode(self, encoding: np.ndarray, x: np.ndarray, t, neighbor_agents_past: np.ndarray) -> Dict[str, Any]: dt = np.full((1, C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, 1), np.float32(t), np.float32) feed = {"encoding": encoding, "sampled_trajectories": np.asarray(x, np.float32).reshape( 1, C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM), "diffusion_time": dt, "neighbor_agents_past": neighbor_agents_past} outs = self.dec.run(["model_output"] + self.decoder_taps, feed) return dict(zip(["model_output"] + self.decoder_taps, outs)) def turn_logit(self, encoding: np.ndarray, final_x0: np.ndarray) -> Dict[str, np.ndarray]: outs = self.turn.run(["turn_indicator_logit"] + self.turn_taps, {"encoding": encoding, "final_x0": np.asarray(final_x0, np.float32)}) return dict(zip(["turn_indicator_logit"] + self.turn_taps, outs)) def run(self, raw: Any, *, steps: int = C.DPM_SOLVER_STEPS, keep_eval_io: bool = False) -> OrtResult: norm = self.normalize(raw) enc = self.encode(norm) encoding = enc["encoding"] taps: Dict[str, np.ndarray] = {f"enc:{k}": v for k, v in enc.items() if k != "encoding"} cs = decoder_masks(norm).current_states nb = norm["neighbor_agents_past"] eval_in: List[np.ndarray] = [] eval_out: List[np.ndarray] = [] def model_fn(x: np.ndarray, t) -> np.ndarray: k = len(eval_out) out = self.decode(encoding, x, t, nb) for name in self.decoder_taps: taps[f"dec{k}:{name}"] = out[name] y = out["model_output"][0] eval_out.append(y if keep_eval_io else np.empty(0)) if keep_eval_io: eval_in.append(np.array(x, np.float32, copy=True)) return y res = dpm_solver_sample(norm["sampled_trajectories"][0], model_fn, lambda x: apply_prefix_constraint(x, cs), steps) tl = self.turn_logit(encoding, res.final_x[None]) taps.update({f"turn:{k}": v for k, v in tl.items() if k != "turn_indicator_logit"}) return OrtResult(norm, encoding, res.final_x[None], tl["turn_indicator_logit"], res.denoising_steps, res.denoising_timesteps, res.eval_times, eval_in, eval_out if keep_eval_io else [], taps)