| |
| """ONNX Runtime golden: the deployed ``frnet_<variant>.onnx`` of AutowareFoundation/lidar_frnet@v2.0 run as-is (CPU |
| execution provider; re-runs are bit-identical, S:frnet:332) on the four unpadded network inputs, exactly as |
| the node binds them (``points`` float32 (N, 4), ``coors`` int64 (N, 3), ``voxel_coors`` int64 (M, 3), |
| ``inverse_map`` int64 (N,)). |
| |
| With ``intermediates=True`` the in-memory graph also exposes the ONNX tensors that hold the same values as the torch |
| reference's taps (the mapping is ``FRNetWeights.onnx_taps``, recovered from the graph by ``reference/weights.py``), |
| so ``tests/test_reference_cpu.py`` compares module by module. Adding graph outputs changes nothing that is |
| computed. Choose which taps to expose with ``tap_patterns`` (fnmatch); OT128 module taps at 146k points are ~1 GB. |
| |
| onnxruntime is a research / test dependency only (not in the image): it is imported lazily, and the tests skip |
| where it is absent. |
| """ |
| from __future__ import annotations |
|
|
| import fnmatch |
| from typing import Any, Dict, Optional, Sequence |
|
|
| import numpy as np |
|
|
| from ..ttaw.golden import NULL_TAPS, TapRegistry |
| from .config import variant_files |
| from .weights import MODULE_TAPS, FRNetWeights, load_weights |
|
|
| __all__ = ["OrtFRNet", "onnxruntime_available"] |
|
|
|
|
| def onnxruntime_available() -> bool: |
| import importlib.util |
|
|
| return importlib.util.find_spec("onnxruntime") is not None and importlib.util.find_spec("onnx") is not None |
|
|
|
|
| def _session(model_bytes: bytes, threads: int): |
| import onnxruntime as ort |
|
|
| so = ort.SessionOptions() |
| so.intra_op_num_threads = int(threads) |
| so.inter_op_num_threads = 1 |
| so.log_severity_level = 3 |
| return ort.InferenceSession(model_bytes, so, providers=["CPUExecutionProvider"]) |
|
|
|
|
| def _with_outputs(path: Any, tensors: Sequence[str]) -> bytes: |
| """The ONNX model at ``path`` with ``tensors`` added as graph outputs (in memory).""" |
| import onnx |
| from onnx import TensorProto, helper |
|
|
| model = onnx.load(str(path)) |
| existing = {o.name for o in model.graph.output} |
| for t in tensors: |
| if t not in existing: |
| model.graph.output.append(helper.make_tensor_value_info(t, TensorProto.FLOAT, None)) |
| existing.add(t) |
| return model.SerializeToString() |
|
|
|
|
| class OrtFRNet: |
| """The deployed graph of one variant under ONNX Runtime (module docstring).""" |
|
|
| def __init__(self, weights_dir: Any, variant: str, *, threads: int = 2, intermediates: bool = False, |
| tap_patterns: Sequence[str] = MODULE_TAPS, weights: Optional[FRNetWeights] = None): |
| self.variant = variant |
| self.path = variant_files(weights_dir, variant)["onnx"] |
| self.weights = weights or load_weights(weights_dir, variant) |
| self.taps: Dict[str, str] = {} |
| if intermediates: |
| self.taps = {name: tensor for name, tensor in self.weights.onnx_taps.items() |
| if name != "pred_probs" and any(fnmatch.fnmatchcase(name, p) for p in tap_patterns)} |
| model_bytes = _with_outputs(self.path, list(self.taps.values())) |
| else: |
| model_bytes = self.path.read_bytes() |
| self.session = _session(model_bytes, threads) |
| self.output_names = [o.name for o in self.session.get_outputs()] |
|
|
| def __call__(self, *args: Any, **kwargs: Any) -> np.ndarray: |
| return self.forward(*args, **kwargs) |
|
|
| def forward(self, points: np.ndarray, coors: np.ndarray, voxel_coors: np.ndarray, inverse_map: np.ndarray, |
| taps: TapRegistry = NULL_TAPS) -> np.ndarray: |
| feeds = {"points": np.ascontiguousarray(points, np.float32), "coors": np.ascontiguousarray(coors, np.int64), |
| "voxel_coors": np.ascontiguousarray(voxel_coors, np.int64), |
| "inverse_map": np.ascontiguousarray(inverse_map, np.int64)} |
| outs = dict(zip(self.output_names, self.session.run(None, feeds))) |
| for name, tensor in self.taps.items(): |
| if taps.wants(name): |
| taps.tap(name, outs[tensor]) |
| probs = outs["pred_probs"] |
| taps.tap("pred_probs", probs) |
| return probs |
|
|
