File size: 4,275 Bytes
5e23710 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | # coding: utf-8
"""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]
|