File size: 5,296 Bytes
dc79b9f | 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 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | #!/usr/bin/env python3
from __future__ import annotations
import json
import wave
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(int16)", "int16", "s16"), np.int16),
(("tensor(int8)", "int8", "s8"), np.int8),
(("tensor(uint16)", "uint16", "u16"), np.uint16),
(("tensor(uint8)", "uint8", "u8"), np.uint8),
(("tensor(bool)", "bool"), np.bool_),
)
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)
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:
"""Small common wrapper for ONNX Runtime and AXEngine."""
def __init__(self, model_path: str | Path, backend: str):
self.path = Path(model_path)
if not self.path.is_file():
raise FileNotFoundError(self.path)
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]
def read_wav(path: str | Path) -> tuple[np.ndarray, int]:
path = Path(path)
with wave.open(str(path), "rb") as source:
if source.getnchannels() != 1 or source.getsampwidth() != 2:
raise ValueError(f"expected mono 16-bit PCM WAV: {path}")
sample_rate = source.getframerate()
samples = np.frombuffer(
source.readframes(source.getnframes()), dtype=np.int16
).copy()
return samples, sample_rate
def load_json(path: str | Path) -> dict[str, Any]:
return json.loads(Path(path).read_text())
def write_json(path: str | Path, value: Any) -> None:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n")
def max_abs_delta(reference: Any, candidate: Any) -> float:
reference_array = np.asarray(reference, dtype=np.float64)
candidate_array = np.asarray(candidate, dtype=np.float64)
if reference_array.shape != candidate_array.shape:
return float("inf")
if reference_array.size == 0:
return 0.0
return float(np.max(np.abs(reference_array - candidate_array)))
|