| |
| """axengine / onnxruntime 推理会话封装。""" |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
|
|
|
|
| @dataclass(frozen=True) |
| class TensorInfo: |
| name: str |
| shape: tuple[int, ...] |
| dtype: np.dtype |
|
|
|
|
| def _numpy_dtype(value: Any) -> np.dtype: |
| text = str(value).lower() |
| mapping = ( |
| (("tensor(float)", "float32", "fp32", "f32"), np.float32), |
| (("tensor(float16)", "float16", "fp16", "f16"), np.float16), |
| (("tensor(int64)", "int64", "s64"), np.int64), |
| (("tensor(int32)", "int32", "s32"), np.int32), |
| (("tensor(uint16)", "uint16", "u16"), np.uint16), |
| (("tensor(uint8)", "uint8", "u8"), np.uint8), |
| ) |
| for aliases, dtype in mapping: |
| if any(alias in text for alias in aliases): |
| return np.dtype(dtype) |
| raise ValueError(f"unsupported runtime tensor dtype: {value}") |
|
|
|
|
| def _tensor_info(value: Any) -> TensorInfo: |
| shape = getattr(value, "shape", None) |
| if shape is None: |
| shape = getattr(value, "dims", None) |
| if shape is None or any(dim is None for dim in shape): |
| raise ValueError(f"dynamic or missing tensor shape for {value.name}: {shape}") |
| dtype = getattr(value, "dtype", None) |
| if dtype is None: |
| dtype = getattr(value, "type", None) |
| return TensorInfo( |
| name=value.name, |
| shape=tuple(int(dim) for dim in shape), |
| dtype=_numpy_dtype(dtype), |
| ) |
|
|
|
|
| class InferenceSession: |
| """.axmodel -> axengine；.onnx -> onnxruntime CPU。""" |
|
|
| def __init__(self, model_path: str | Path, backend: str | None = None): |
| self.path = Path(model_path) |
| if not self.path.is_file(): |
| raise FileNotFoundError(self.path) |
| if backend is None: |
| backend = "axengine" if self.path.suffix == ".axmodel" else "onnx" |
| self.backend = backend |
| if backend == "axengine": |
| try: |
| import axengine |
| except ImportError as error: |
| raise RuntimeError( |
| "axengine is unavailable; run this backend on an AXERA board" |
| ) from error |
| self._session = axengine.InferenceSession(str(self.path)) |
| elif backend == "onnx": |
| try: |
| import onnxruntime as ort |
| except ImportError as error: |
| raise RuntimeError("onnxruntime is required for --backend onnx") from error |
| options = ort.SessionOptions() |
| options.inter_op_num_threads = 1 |
| options.intra_op_num_threads = 1 |
| self._session = ort.InferenceSession( |
| str(self.path), |
| sess_options=options, |
| providers=["CPUExecutionProvider"], |
| ) |
| else: |
| raise ValueError(f"unsupported backend: {backend}") |
|
|
| self.inputs = [_tensor_info(value) for value in self._session.get_inputs()] |
| self.outputs = [_tensor_info(value) for value in self._session.get_outputs()] |
| self.input_by_name = {value.name: value for value in self.inputs} |
|
|
| def run(self, feed: dict[str, np.ndarray]) -> dict[str, np.ndarray]: |
| missing = [value.name for value in self.inputs if value.name not in feed] |
| if missing: |
| raise KeyError(f"missing inputs for {self.path.name}: {missing}") |
| prepared = { |
| name: np.ascontiguousarray(np.asarray(feed[name], dtype=meta.dtype)) |
| for name, meta in self.input_by_name.items() |
| } |
| values = self._session.run(None, prepared) |
| if isinstance(values, dict): |
| return {name: np.asarray(value) for name, value in values.items()} |
| if not isinstance(values, (list, tuple)): |
| values = [values] |
| if len(values) != len(self.outputs): |
| raise RuntimeError( |
| f"unexpected output count from {self.path.name}: " |
| f"{len(values)} != {len(self.outputs)}" |
| ) |
| return { |
| meta.name: np.asarray(value) |
| for meta, value in zip(self.outputs, values) |
| } |
|
|
|
|
| def first_output(session: InferenceSession, outputs: dict[str, np.ndarray]) -> np.ndarray: |
| return outputs[session.outputs[0].name] |
|
|
