File size: 3,621 Bytes
5eee449 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """Inference backend wrapper: pyaxengine (default) or onnxruntime (host stand-in).
Both expose an onnxruntime-like API (run / get_inputs / get_outputs). The
backend is selected from the model file suffix unless overridden:
*.axmodel -> axengine.InferenceSession (pyaxengine, on-device or host sim)
*.onnx -> onnxruntime.InferenceSession (host numeric stand-in only)
Encoder input dtype note: the compiled AXMODELs declare their token inputs as
S32 at runtime (COMPILE_NOTES §3 folds the S64 input_processors into the
model; SIMULATE §1 confirmed S32 binaries), while the ONNX graphs take int64.
The wrapper therefore adapts the feed dtype from the session's input metadata.
"""
from __future__ import annotations
from pathlib import Path
import numpy as np
class ModelSession:
"""Minimal name-addressable inference session over axengine/onnxruntime."""
def __init__(self, model_path: str | Path, backend: str = "auto") -> None:
self.path = Path(model_path)
if not self.path.is_file():
raise FileNotFoundError(self.path)
if backend == "auto":
backend = "axengine" if self.path.suffix == ".axmodel" else "onnxruntime"
self.backend = backend
if backend == "axengine":
try:
from axengine import InferenceSession
except ImportError as exc:
raise ImportError(
"pyaxengine is required for .axmodel files "
"(pip install axengine on an AXera host/BSP environment)"
) from exc
self._sess = InferenceSession(str(self.path))
elif backend == "onnxruntime":
try:
import onnxruntime as ort
except ImportError as exc:
raise ImportError(
"onnxruntime is required for .onnx files "
"(host numeric stand-in only)"
) from exc
self._sess = ort.InferenceSession(
str(self.path), providers=["CPUExecutionProvider"]
)
else:
raise ValueError(f"unknown backend {backend!r}")
# -- metadata -----------------------------------------------------------
def _inputs(self):
return self._sess.get_inputs()
def output_names(self) -> list[str]:
return [o.name for o in self._sess.get_outputs()]
def input_dtype(self, name: str) -> np.dtype:
"""Best-effort numpy dtype of a named input.
onnxruntime exposes `.type` strings ('tensor(int64)'); pyaxengine
exposes numpy-like `.dtype`. Fallback: int32 for .axmodel (see module
docstring), int64 for .onnx.
"""
for meta in self._inputs():
if meta.name != name:
continue
dt = getattr(meta, "dtype", None)
if dt is not None:
try:
return np.dtype(dt)
except TypeError:
pass
type_str = str(getattr(meta, "type", "") or "")
if "int32" in type_str:
return np.dtype(np.int32)
if "int64" in type_str:
return np.dtype(np.int64)
if "float" in type_str:
return np.dtype(np.float32)
return np.dtype(np.int32 if self.path.suffix == ".axmodel" else np.int64)
# -- inference ----------------------------------------------------------
def run(self, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
results = self._sess.run(None, feeds)
return dict(zip(self.output_names(), results))
|