| import numpy as np | |
| DEFAULT_PROVIDER = "AxEngineExecutionProvider" | |
| class ModelSession: | |
| """NPU 专用推理会话：仅支持 AX 芯片端到端运行（pyaxengine），无 CPU/onnxruntime 回退。 | |
| input/output 名称与 export/model_meta.json 一致（AXMODEL 即按此编译）。 | |
| 非 AX 环境（缺少 pyaxengine 或 AX provider）时直接报错，不做 CPU 兜底。 | |
| """ | |
| def __init__(self, model_path, providers=None): | |
| try: | |
| import axengine as axe | |
| except ImportError as exc: | |
| raise RuntimeError( | |
| "SDK 为 NPU 专用发布版，仅支持在 AX 芯片上运行；请先安装 requirements.txt " | |
| "并在板端执行（无 onnxruntime/torch/transformers 回退）" | |
| ) from exc | |
| self.session = axe.InferenceSession( | |
| model_path, providers=providers or [DEFAULT_PROVIDER]) | |
| self.backend = "axengine" | |
| self.input_names = [i.name for i in self.session.get_inputs()] | |
| self.output_names = [o.name for o in self.session.get_outputs()] | |
| def run_named(self, feeds, names=None): | |
| """feeds: 与 names（默认 input_names）对应的数组列表。返回输出列表。""" | |
| names = names or self.input_names | |
| if len(feeds) != len(names): | |
| raise ValueError(f"输入数量不匹配: {len(feeds)} != {len(names)}") | |
| feed = { | |
| name: np.ascontiguousarray(arr, dtype=np.float32) | |
| for name, arr in zip(names, feeds) | |
| } | |
| return self.session.run(None, feed) | |