File size: 4,841 Bytes
26f00f8 | 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 | """3 个 axmodel 的 axengine 会话封装 (enc/erb_dec/df_dec 流水线)。"""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
T = 99 # 静态时间帧 (1s @ 48kHz 流式 STFT)
@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 _AxSession:
"""单个 axmodel 的 axengine 会话 (参考 AXERA 官方推理工程封装)。"""
def __init__(self, model_path: str | Path):
import axengine
self.path = Path(model_path)
if not self.path.is_file():
raise FileNotFoundError(self.path)
self._session = axengine.InferenceSession(str(self.path))
self.inputs = [_tensor_info(v) for v in self._session.get_inputs()]
self.outputs = [_tensor_info(v) for v in self._session.get_outputs()]
self.input_by_name = {v.name: v for v in self.inputs}
def run(self, feed: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
missing = [v.name for v in self.inputs if v.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(v) for name, v 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(v) for meta, v in zip(self.outputs, values)}
class DfNet3Ax:
"""DeepFilterNet3 三子图 NPU 流水线 (与 dsp.pipeline 配合)。
enc: feat_erb(1,1,T,32)+feat_spec(1,2,T,96) -> e0..e3, emb, c0, lsnr
erb_dec: emb,e3,e2,e1,e0 -> m(1,1,T,32)
df_dec: emb,c0 -> coefs(1,T,96,10)
"""
def __init__(self, model_dir: str | Path):
model_dir = Path(model_dir)
self.enc = _AxSession(model_dir / "enc.axmodel")
self.erb_dec = _AxSession(model_dir / "erb_dec.axmodel")
self.df_dec = _AxSession(model_dir / "df_dec.axmodel")
# 状态携带模式检测: axmodel 是否有 h0 输入
self.stateful = all("h0" in s.input_by_name for s in (self.enc, self.df_dec))
def run_enc(self, feat_erb: np.ndarray, feat_spec: np.ndarray,
h: np.ndarray | None = None) -> tuple[dict[str, np.ndarray], np.ndarray | None]:
feed = {"feat_erb": feat_erb, "feat_spec": feat_spec}
if self.stateful:
if h is None:
h = np.zeros((1, 1, 256), np.float32)
feed["h0"] = h
out = self.enc.run(feed)
return out, np.asarray(out["h_out"])
return self.enc.run(feed), None
def run_erb_dec(self, emb, e3, e2, e1, e0) -> np.ndarray:
return self.erb_dec.run({"emb": emb, "e3": e3, "e2": e2, "e1": e1, "e0": e0})["m"]
def run_df_dec(self, emb, c0, h: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray | None]:
feed = {"emb": emb, "c0": c0}
if self.stateful:
if h is None:
h = np.zeros((2, 1, 256), np.float32)
feed["h0"] = h
out = self.df_dec.run(feed)
return np.asarray(out["coefs"]), np.asarray(out["h_out"])
return self.df_dec.run(feed)["coefs"], None
|