| |
| |
| """ |
| 把 DepthART ONNX 中的 com.depthart::SelectiveScan 改写为 pulsar2 可编译的纯 ONNX 子图。 |
| |
| 背景： |
| 官方 relative_* ONNX 含自定义算子 SelectiveScan，pulsar2 前端直接报 |
| "dont support SelectiveScan opr"。官方只提供 TensorRT 插件，没有 NPU 实现。 |
| 推理脚本里的 Loop 改写可以跑通 onnxruntime，但 AX 工具链不支持 Loop 控制流。 |
| |
| 做法： |
| 选择性扫描是线性递推 h_t = exp(dt_t*A) * h_{t-1} + dt_t*B_t*u_t，可写成 |
| h_t = exp(L_t) * h_{-1} + sum_{s<=t} exp(L_t - L_s) * dt_s*B_s*u_s |
| 其中 L_t = cumsum_{r<=t}(dt_r*A)。把时间轴切成固定长度 T（默认 14），每个 |
| 分块内用 CumSum + 下三角 mask + MatMul 一次算出，块间只传一个状态。所有 |
| 指数项的指数都 <= 0，不会溢出。所有算子（CumSum/Exp/Where/MatMul 等）都在 |
| AX 前端支持列表内，不含 Loop/Scan/自定义算子。 |
| |
| 校验： |
| - 本工具用 Identity 替换 scan 跑一遍，拿到每个 scan 的 D/G/N/L，再按实际 |
| 尺寸生成静态图；时间轴不是 T 整数倍时自动补零并在输出处裁掉。 |
| - 改写后的整模型输出与已验证的 Loop 版本最大绝对误差约 6e-6（FP32）。 |
| - 同时把 4D `Einsum` 换成按 k 的批 MatMul（`b k d l, k c d -> b k c l` |
| 和 `b k r l, k d r -> b k d l`）：pulsar2 的量化执行器不支持 4D Einsum， |
| 会报 Op Execution Error。 |
| |
| 用法（在 hs/DepthART 目录下执行）： |
| # 默认：448 模型，块大小 14，输出 onnx/relative_s_448_default_axonnx.onnx |
| ../../venv/bin/python src/tool/depthart_export_axonnx.py |
| |
| # 指定模型/输出/块大小 |
| ../../venv/bin/python src/tool/depthart_export_axonnx.py \\ |
| --model onnx/relative_b_448_default.onnx \\ |
| --output onnx/relative_b_448_default_axonnx.onnx --chunk-size 14 |
| |
| # 改写后用推理脚本验证数值 |
| ../../venv/bin/python src/infer/depthart_infer_onnx.py \\ |
| --model onnx/relative_s_448_default_axonnx.onnx --image /path/a.jpg |
| |
| 注意： |
| - 块内两两指数矩阵形状为 [D*N, T, T]，T 越大显存/内存占用越大； |
| T=14 时 672 通道层约 1M 元素（4 MB），推荐保持默认。 |
| - 输出是标准 ONNX，不再含自定义算子；原模型文件不会被修改。 |
| """ |
|
|
| import argparse |
| import logging |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnx |
| import onnxruntime as ort |
| from onnx import helper, numpy_helper, TensorProto |
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[2] |
| DEFAULT_MODEL = PROJECT_ROOT / 'onnx' / 'relative_s_448_default.onnx' |
| DOMAIN = 'com.depthart' |
| OP = 'SelectiveScan' |
|
|
|
|
| EINSUM_EQUATIONS = { |
| 'b k d l, k c d -> b k c l', |
| 'b k r l, k d r -> b k d l', |
| } |
|
|
|
|
| def collect_runtime_shapes(model): |
| """用 ORT 实际跑一遍，读取 scan 与 Einsum 的静态形状。 |
| |
| scan 会被替换成 Identity（输出形状等于 u，不影响下游形状）；Einsum 正常执行。 |
| 返回 (scan_shapes, einsum_shapes)： |
| scan_shapes : 按图顺序的 [(D, G, N, L), ...] |
| einsum_shapes: {节点名: (x_shape, y_shape)} |
| """ |
| nodes = [] |
| scan_outputs, b_aliases = [], {} |
| einsum_x, einsum_y = {}, {} |
| for index, node in enumerate(model.graph.node): |
| if node.domain == DOMAIN and node.op_type == OP: |
| nodes.append(helper.make_node('Identity', [node.input[0]], [node.output[0]], |
| name=f'depthart_stub_scan_{index}')) |
| alias = f'depthart_stub_b_{index}' |
| nodes.append(helper.make_node('Identity', [node.input[3]], [alias], |
| name=f'depthart_stub_bnode_{index}')) |
| scan_outputs.append(node.output[0]) |
| b_aliases[node.output[0]] = alias |
| else: |
| if node.op_type == 'Einsum': |
| x_alias = f'depthart_stub_einsum_x_{index}' |
| y_alias = f'depthart_stub_einsum_y_{index}' |
| nodes.append(helper.make_node('Identity', [node.input[0]], [x_alias], |
| name=f'{x_alias}_node')) |
| nodes.append(helper.make_node('Identity', [node.output[0]], [y_alias], |
| name=f'{y_alias}_node')) |
| einsum_x[node.name] = x_alias |
| einsum_y[node.name] = y_alias |
| nodes.append(node) |
|
|
| stub = onnx.ModelProto() |
| stub.CopyFrom(model) |
| del stub.graph.node[:] |
| stub.graph.node.extend(nodes) |
| for name in scan_outputs: |
| stub.graph.output.append(helper.make_tensor_value_info(name, TensorProto.FLOAT, None)) |
| stub.graph.output.append(helper.make_tensor_value_info(b_aliases[name], |
| TensorProto.FLOAT, None)) |
| for x_alias, y_alias in zip(einsum_x.values(), einsum_y.values()): |
| stub.graph.output.append(helper.make_tensor_value_info(x_alias, TensorProto.FLOAT, None)) |
| stub.graph.output.append(helper.make_tensor_value_info(y_alias, TensorProto.FLOAT, None)) |
|
|
| meta = model.graph.input[0] |
| shape = [dim.dim_value for dim in meta.type.tensor_type.shape.dim] |
| if not all(shape): |
| raise ValueError(f'输入必须是静态形状: {meta.name} {shape}') |
| session = ort.InferenceSession(stub.SerializeToString()) |
| values = dict(zip([out.name for out in session.get_outputs()], |
| session.run(None, {meta.name: np.zeros(shape, dtype=np.float32)}))) |
|
|
| scan_map = {} |
| for name in scan_outputs: |
| u_shape = values[name].shape |
| b_shape = values[b_aliases[name]].shape |
| if len(u_shape) != 3 or len(b_shape) != 4: |
| raise ValueError(f'scan 形状不符合预期: u={u_shape} B={b_shape}') |
| scan_map[name] = (int(u_shape[1]), int(b_shape[1]), int(b_shape[2]), int(u_shape[2])) |
| scan_shapes = [scan_map[node.output[0]] for node in model.graph.node |
| if node.domain == DOMAIN and node.op_type == OP] |
| einsum_shapes = {name: (values[einsum_x[name]].shape, values[einsum_y[name]].shape) |
| for name in einsum_x} |
| return scan_shapes, einsum_shapes |
|
|
|
|
| class Rewriter: |
| def __init__(self, model, chunk_size, stub_scan=False, scan_method='pairwise'): |
| self.model = model |
| self.chunk_size = chunk_size |
| self.stub_scan = stub_scan |
| self.scan_method = scan_method |
| self.initializers = list(model.graph.initializer) |
| self.names = {tensor.name for tensor in self.initializers} |
| self.initializer_shapes = {tensor.name: tuple(tensor.dims) |
| for tensor in model.graph.initializer} |
| self.shared = {} |
|
|
| def build_einsum(self, prefix, node, x_shape, y_shape): |
| """b k d l, k c d -> b k c l （以及 b k r l, k d r -> b k d l）。 |
| |
| 两条式子都是按 k 的批 MatMul：W[k,out,in] @ x[b=1,k,in,l]。 |
| 替换掉 4D Einsum，避免量化执行器不支持。 |
| """ |
| out_name = node.output[0] |
| w_shape = self.initializer_shapes.get(node.input[1]) |
| if w_shape is None: |
| raise ValueError(f'{node.name}: 第二个输入不是常量权重，暂不支持') |
| if len(w_shape) != 3 or len(x_shape) != 4: |
| raise ValueError(f'{node.name}: Einsum 形状不符合预期: {w_shape} {x_shape}') |
| if int(x_shape[0]) != 1: |
| raise ValueError(f'{node.name}: 仅支持 batch=1，实际 {x_shape}') |
| K, out_dim, in_dim = (int(v) for v in w_shape) |
| length = int(x_shape[3]) |
| if int(x_shape[1]) != K: |
| raise ValueError(f'{node.name}: K 不一致: {x_shape} vs {w_shape}') |
| expected = (1, K, out_dim, length) |
| if tuple(int(v) for v in y_shape) != expected: |
| raise ValueError(f'{node.name}: 输出形状 {y_shape} != {expected}') |
| nodes = [] |
| |
| |
| nodes.append(helper.make_node( |
| 'Squeeze', [node.input[0], self.shared_const('axes0', np.array([0], np.int64))], |
| [f'{prefix}_x2'])) |
| nodes.append(helper.make_node('MatMul', [node.input[1], f'{prefix}_x2'], |
| [f'{prefix}_y3'])) |
| nodes.append(helper.make_node( |
| 'Unsqueeze', [f'{prefix}_y3', self.shared_const('axes0', np.array([0], np.int64))], |
| [out_name])) |
| return nodes |
|
|
| def const(self, name, value): |
| if name not in self.names: |
| self.initializers.append(numpy_helper.from_array(np.asarray(value), name)) |
| self.names.add(name) |
| return name |
|
|
| def shared_const(self, key, value): |
| if key not in self.shared: |
| self.shared[key] = self.const(f'depthart_{key}', value) |
| return self.shared[key] |
|
|
| def build_logspace(self, prefix, shapes, inputs): |
| """对数空间分块递推：每块 O(T) 而非 O(T^2)，不构造 [D*N,T,T] 两两矩阵。 |
| |
| h_t = exp(L_t) * ( h_prev + cumsum_s( exp(-L_s) * dt_s * B_s * u_s ) ) |
| 其中 L 为块内相对累计衰减（L_0=0，L_t<=0）。要求块内 exp(-L) 不超 U16 量程， |
| 因此 T 要小（T=4 时最大约 e^9≈8e3）。 |
| """ |
| D, G, N, L = shapes |
| u, delta, a, b, c = inputs[:5] |
| d = inputs[5] if len(inputs) > 5 and inputs[5] else None |
| bias = inputs[6] if len(inputs) > 6 and inputs[6] else None |
| attrs = {item.name: item.i for item in self.node_attrs} |
| T = self.chunk_size |
| nodes = [] |
|
|
| dt = delta |
| if bias is not None: |
| bias3 = self.const(f'{prefix}_bias3_shape', np.array([1, D, 1], np.int64)) |
| nodes.append(helper.make_node('Reshape', [bias, bias3], [f'{prefix}_bias3'])) |
| nodes.append(helper.make_node('Add', [delta, f'{prefix}_bias3'], [f'{prefix}_dtb'])) |
| dt = f'{prefix}_dtb' |
| if int(attrs.get('delta_softplus', 0)): |
| nodes.append(helper.make_node('Softplus', [dt], [f'{prefix}_dt'])) |
| dt = f'{prefix}_dt' |
|
|
| |
| pad = (-L) % T |
| Lp = L + pad |
| if pad: |
| nodes.append(helper.make_node( |
| 'Concat', [dt, self.const(f'{prefix}_pad_dt', np.zeros((1, D, pad), np.float32))], |
| [f'{prefix}_dt_pad'], axis=2)) |
| nodes.append(helper.make_node( |
| 'Concat', [u, self.const(f'{prefix}_pad_u', np.zeros((1, D, pad), np.float32))], |
| [f'{prefix}_u_pad'], axis=2)) |
| pad_bc = self.const(f'{prefix}_pad_bc', np.zeros((1, G, N, pad), np.float32)) |
| nodes.append(helper.make_node('Concat', [b, pad_bc], [f'{prefix}_b_pad'], axis=3)) |
| nodes.append(helper.make_node('Concat', [c, pad_bc], [f'{prefix}_c_pad'], axis=3)) |
| dt, u, b, c = (f'{prefix}_dt_pad', f'{prefix}_u_pad', |
| f'{prefix}_b_pad', f'{prefix}_c_pad') |
|
|
| |
| a4 = self.const(f'{prefix}_a4_shape', np.array([1, D, N, 1], np.int64)) |
| nodes.append(helper.make_node('Reshape', [a, a4], [f'{prefix}_a4'])) |
| dt3 = self.const(f'{prefix}_dt3_shape', np.array([1, D, 1, Lp], np.int64)) |
| nodes.append(helper.make_node('Reshape', [dt, dt3], [f'{prefix}_dt4'])) |
| nodes.append(helper.make_node('Mul', [f'{prefix}_dt4', f'{prefix}_a4'], |
| [f'{prefix}_dtA'])) |
|
|
| b_full = self.expand_groups(b, D, G, N, Lp, f'{prefix}_b', nodes) |
| c_full = self.expand_groups(c, D, G, N, Lp, f'{prefix}_c', nodes) |
|
|
| if d is not None: |
| nodes.append(helper.make_node('Reshape', [d, |
| self.const(f'{prefix}_d_shape', |
| np.array([1, D, 1], np.int64))], |
| [f'{prefix}_d3'])) |
| d = f'{prefix}_d3' |
|
|
| axes_u2 = self.shared_const('axes_u2', np.array([2], np.int64)) |
| axes_rs = self.shared_const('axes_rs2', np.array([2], np.int64)) |
| axis_t = self.shared_const('axis2', np.array([2], np.int64)) |
| axis_bc = self.shared_const('axis3', np.array([3], np.int64)) |
| h_prev = self.const(f'{prefix}_h0', np.zeros((1, D, N, 1), np.float32)) |
| outputs = [] |
| for k in range(Lp // T): |
| start, end = k * T, (k + 1) * T |
| tag = f'{prefix}_c{k}' |
| starts = self.const(f'{tag}_starts', np.array([start], np.int64)) |
| ends = self.const(f'{tag}_ends', np.array([end], np.int64)) |
| for source, suffix, axis in ((f'{prefix}_dtA', 'dta', axis_bc), |
| (dt, 'dt', axis_t), (u, 'u', axis_t), |
| (b_full, 'b', axis_bc), (c_full, 'c', axis_bc)): |
| nodes.append(helper.make_node('Slice', [source, starts, ends, axis], |
| [f'{tag}_{suffix}'])) |
|
|
| |
| nodes.append(helper.make_node('CumSum', [f'{tag}_dta', |
| self.shared_const('axis3s', np.array(3, np.int64))], |
| [f'{tag}_L'])) |
| nodes.append(helper.make_node('Exp', [f'{tag}_L'], [f'{tag}_eL'])) |
| nodes.append(helper.make_node('Neg', [f'{tag}_L'], [f'{tag}_negL'])) |
| nodes.append(helper.make_node('Exp', [f'{tag}_negL'], [f'{tag}_eneg'])) |
|
|
| |
| nodes.append(helper.make_node('Unsqueeze', [f'{tag}_dt', axes_u2], [f'{tag}_dt3'])) |
| nodes.append(helper.make_node('Unsqueeze', [f'{tag}_u', axes_u2], [f'{tag}_u3'])) |
| nodes.append(helper.make_node('Mul', [f'{tag}_dt3', f'{tag}_u3'], [f'{tag}_dtu'])) |
| nodes.append(helper.make_node('Mul', [f'{tag}_dtu', f'{tag}_b'], [f'{tag}_bt'])) |
| nodes.append(helper.make_node('Mul', [f'{tag}_eneg', f'{tag}_bt'], [f'{tag}_g'])) |
| nodes.append(helper.make_node('CumSum', [f'{tag}_g', |
| self.shared_const('axis3s', np.array(3, np.int64))], |
| [f'{tag}_G'])) |
| nodes.append(helper.make_node('Add', [f'{tag}_G', h_prev], [f'{tag}_hp'])) |
| nodes.append(helper.make_node('Mul', [f'{tag}_eL', f'{tag}_hp'], [f'{tag}_h'])) |
|
|
| nodes.append(helper.make_node('Mul', [f'{tag}_h', f'{tag}_c'], [f'{tag}_hc'])) |
| nodes.append(helper.make_node('ReduceSum', [f'{tag}_hc', axes_rs], |
| [f'{tag}_ysum'], keepdims=0)) |
| if d is not None: |
| nodes.append(helper.make_node('Mul', [d, f'{tag}_u'], [f'{tag}_du'])) |
| nodes.append(helper.make_node('Add', [f'{tag}_ysum', f'{tag}_du'], |
| [f'{tag}_y'])) |
| else: |
| nodes.append(helper.make_node('Identity', [f'{tag}_ysum'], [f'{tag}_y'])) |
| outputs.append(f'{tag}_y') |
|
|
| nodes.append(helper.make_node('Slice', [f'{tag}_h', |
| self.const(f'{tag}_hl', np.array([T - 1], np.int64)), |
| self.const(f'{tag}_he', np.array([T], np.int64)), |
| axis_bc], [f'{tag}_hnext'])) |
| h_prev = f'{tag}_hnext' |
|
|
| nodes.append(helper.make_node('Concat', outputs, [f'{prefix}_ycat'], axis=2)) |
| if pad: |
| nodes.append(helper.make_node('Slice', [ |
| f'{prefix}_ycat', |
| self.const(f'{prefix}_out_start', np.array([0], np.int64)), |
| self.const(f'{prefix}_out_end', np.array([L], np.int64)), |
| self.shared_const('axis2', np.array([2], np.int64))], [self.out_name])) |
| else: |
| nodes.append(helper.make_node('Identity', [f'{prefix}_ycat'], [self.out_name])) |
| return nodes |
|
|
| def build(self, prefix, shapes, inputs): |
| D, G, N, L = shapes |
| u, delta, a, b, c = inputs[:5] |
| d = inputs[5] if len(inputs) > 5 and inputs[5] else None |
| bias = inputs[6] if len(inputs) > 6 and inputs[6] else None |
| attrs = {item.name: item.i for item in self.node_attrs} |
| T = self.chunk_size |
| nodes = [] |
|
|
| |
| dt = delta |
| if bias is not None: |
| bias3 = self.const(f'{prefix}_bias3_shape', np.array([1, D, 1], np.int64)) |
| nodes.append(helper.make_node('Reshape', [bias, bias3], [f'{prefix}_bias3'])) |
| nodes.append(helper.make_node('Add', [delta, f'{prefix}_bias3'], [f'{prefix}_dtb'])) |
| dt = f'{prefix}_dtb' |
| if int(attrs.get('delta_softplus', 0)): |
| nodes.append(helper.make_node('Softplus', [dt], [f'{prefix}_dt'])) |
| dt = f'{prefix}_dt' |
|
|
| |
| pad = (-L) % T |
| Lp = L + pad |
| if pad: |
| nodes.append(helper.make_node( |
| 'Concat', [dt, self.const(f'{prefix}_pad_dt', np.zeros((1, D, pad), np.float32))], |
| [f'{prefix}_dt_pad'], axis=2)) |
| nodes.append(helper.make_node( |
| 'Concat', [u, self.const(f'{prefix}_pad_u', np.zeros((1, D, pad), np.float32))], |
| [f'{prefix}_u_pad'], axis=2)) |
| pad_bc = self.const(f'{prefix}_pad_bc', np.zeros((1, G, N, pad), np.float32)) |
| nodes.append(helper.make_node('Concat', [b, pad_bc], [f'{prefix}_b_pad'], axis=3)) |
| nodes.append(helper.make_node('Concat', [c, pad_bc], [f'{prefix}_c_pad'], axis=3)) |
| dt, u, b, c = (f'{prefix}_dt_pad', f'{prefix}_u_pad', |
| f'{prefix}_b_pad', f'{prefix}_c_pad') |
|
|
| |
| b_full = self.expand_groups(b, D, G, N, Lp, f'{prefix}_b', nodes) |
| c_full = self.expand_groups(c, D, G, N, Lp, f'{prefix}_c', nodes) |
|
|
| |
| a4 = self.const(f'{prefix}_a_shape', np.array([1, D, N, 1], np.int64)) |
| nodes.append(helper.make_node('Reshape', [a, a4], [f'{prefix}_a4'])) |
| if d is not None: |
| d3 = self.const(f'{prefix}_d_shape', np.array([1, D, 1], np.int64)) |
| nodes.append(helper.make_node('Reshape', [d, d3], [f'{prefix}_d3'])) |
| d = f'{prefix}_d3' |
|
|
| mask = np.tril(np.ones((T, T), dtype=bool)).reshape(1, T, T) |
| mask_name = self.shared_const(f'mask_{T}', mask) |
| neg_name = self.shared_const('mask_fill', np.array(-1e4, np.float32)) |
| flat_axis = self.shared_const('axis1', np.array(1, np.int64)) |
| axes_u2 = self.shared_const('axes_u2', np.array([2], np.int64)) |
| axes_u3 = self.shared_const('axes_u3', np.array([3], np.int64)) |
| axes_rs = self.shared_const('axes_rs2', np.array([2], np.int64)) |
|
|
| h_prev = self.const(f'{prefix}_h0', np.zeros((1, D, N, 1), np.float32)) |
| outputs = [] |
| for k in range(Lp // T): |
| start, end = k * T, (k + 1) * T |
| tag = f'{prefix}_c{k}' |
| starts = self.const(f'{tag}_starts', np.array([start], np.int64)) |
| ends = self.const(f'{tag}_ends', np.array([end], np.int64)) |
| axis_t = self.shared_const('axis2', np.array([2], np.int64)) |
| axis_bc = self.shared_const('axis3', np.array([3], np.int64)) |
| for source, suffix, axis in ((dt, 'dt', axis_t), (u, 'u', axis_t), |
| (b_full, 'b', axis_bc), (c_full, 'c', axis_bc)): |
| nodes.append(helper.make_node('Slice', [source, starts, ends, axis], |
| [f'{tag}_{suffix}'])) |
|
|
| |
| nodes.append(helper.make_node('Unsqueeze', [f'{tag}_dt', axes_u2], [f'{tag}_dt3'])) |
| nodes.append(helper.make_node('Mul', [f'{tag}_dt3', f'{prefix}_a4'], [f'{tag}_dtA'])) |
| nodes.append(helper.make_node('Reshape', [f'{tag}_dtA', |
| self.const(f'{tag}_dtA2_shape', |
| np.array([D * N, T], np.int64))], |
| [f'{tag}_dtA2'])) |
| nodes.append(helper.make_node('CumSum', [f'{tag}_dtA2', flat_axis], [f'{tag}_L2'])) |
| nodes.append(helper.make_node('Reshape', [f'{tag}_L2', |
| self.const(f'{tag}_L_shape', |
| np.array([1, D, N, T], np.int64))], |
| [f'{tag}_L'])) |
|
|
| |
| nodes.append(helper.make_node('Reshape', [f'{tag}_L2', |
| self.const(f'{tag}_Lt_shape', |
| np.array([D * N, T, 1], np.int64))], |
| [f'{tag}_Lt'])) |
| nodes.append(helper.make_node('Reshape', [f'{tag}_L2', |
| self.const(f'{tag}_Ls_shape', |
| np.array([D * N, 1, T], np.int64))], |
| [f'{tag}_Ls'])) |
| nodes.append(helper.make_node('Sub', [f'{tag}_Lt', f'{tag}_Ls'], [f'{tag}_diff'])) |
| nodes.append(helper.make_node('Where', [mask_name, f'{tag}_diff', neg_name], |
| [f'{tag}_diffm'])) |
| nodes.append(helper.make_node('Exp', [f'{tag}_diffm'], [f'{tag}_E'])) |
|
|
| |
| nodes.append(helper.make_node('Unsqueeze', [f'{tag}_u', axes_u2], [f'{tag}_u3'])) |
| nodes.append(helper.make_node('Mul', [f'{tag}_dt3', f'{tag}_u3'], [f'{tag}_dtu'])) |
| nodes.append(helper.make_node('Mul', [f'{tag}_dtu', f'{tag}_b'], [f'{tag}_bt'])) |
| nodes.append(helper.make_node('Reshape', [f'{tag}_bt', |
| self.const(f'{tag}_bt2_shape', |
| np.array([D * N, 1, T], np.int64))], |
| [f'{tag}_bt2'])) |
| nodes.append(helper.make_node('Mul', [f'{tag}_E', f'{tag}_bt2'], [f'{tag}_eb'])) |
| nodes.append(helper.make_node('ReduceSum', [f'{tag}_eb', |
| self.shared_const('axis2s', np.array([2], np.int64))], |
| [f'{tag}_contrib2'], keepdims=0)) |
| nodes.append(helper.make_node('Reshape', [f'{tag}_contrib2', |
| self.const(f'{tag}_contrib_shape', |
| np.array([1, D, N, T], np.int64))], |
| [f'{tag}_contrib'])) |
|
|
| |
| nodes.append(helper.make_node('Exp', [f'{tag}_L'], [f'{tag}_expL'])) |
| nodes.append(helper.make_node('Mul', [f'{tag}_expL', h_prev], [f'{tag}_carry'])) |
| nodes.append(helper.make_node('Add', [f'{tag}_carry', f'{tag}_contrib'], |
| [f'{tag}_h'])) |
|
|
| |
| nodes.append(helper.make_node('Mul', [f'{tag}_h', f'{tag}_c'], [f'{tag}_hc'])) |
| nodes.append(helper.make_node('ReduceSum', [f'{tag}_hc', axes_rs], |
| [f'{tag}_ysum'], keepdims=0)) |
| if d is not None: |
| nodes.append(helper.make_node('Mul', [d, f'{tag}_u'], [f'{tag}_du'])) |
| nodes.append(helper.make_node('Add', [f'{tag}_ysum', f'{tag}_du'], |
| [f'{tag}_y'])) |
| else: |
| nodes.append(helper.make_node('Identity', [f'{tag}_ysum'], [f'{tag}_y'])) |
| outputs.append(f'{tag}_y') |
|
|
| nodes.append(helper.make_node('Gather', [f'{tag}_h', |
| self.const(f'{tag}_last', np.array(T - 1, np.int64))], |
| [f'{tag}_hlast'], axis=3)) |
| nodes.append(helper.make_node('Unsqueeze', [f'{tag}_hlast', axes_u3], |
| [f'{tag}_hnext'])) |
| h_prev = f'{tag}_hnext' |
|
|
| nodes.append(helper.make_node('Concat', outputs, [f'{prefix}_ycat'], axis=2)) |
| if pad: |
| nodes.append(helper.make_node('Slice', [ |
| f'{prefix}_ycat', |
| self.const(f'{prefix}_out_start', np.array([0], np.int64)), |
| self.const(f'{prefix}_out_end', np.array([L], np.int64)), |
| self.shared_const('axis2', np.array([2], np.int64))], [self.out_name])) |
| else: |
| nodes.append(helper.make_node('Identity', [f'{prefix}_ycat'], [self.out_name])) |
| return nodes |
|
|
| def expand_groups(self, x, D, G, N, L, prefix, nodes): |
| gs = D // G |
| r1_shape = self.const(f'{prefix}_r1_shape', np.array([1, G, 1, N, L], np.int64)) |
| nodes.append(helper.make_node('Reshape', [x, r1_shape], [f'{prefix}_r1'])) |
| target = self.const(f'{prefix}_tgt', np.array([1, G, gs, N, L], np.int64)) |
| nodes.append(helper.make_node('Expand', [f'{prefix}_r1', target], [f'{prefix}_e'])) |
| r2_shape = self.const(f'{prefix}_r2_shape', np.array([1, D, N, L], np.int64)) |
| nodes.append(helper.make_node('Reshape', [f'{prefix}_e', r2_shape], [f'{prefix}_bc4'])) |
| return f'{prefix}_bc4' |
|
|
| def rewrite(self): |
| shapes, einsum_shapes = collect_runtime_shapes(self.model) |
| if not shapes: |
| raise ValueError('模型中没有 SelectiveScan 节点，无需改写') |
| new_nodes, index, einsum_index = [], 0, 0 |
| for node in self.model.graph.node: |
| if node.domain == DOMAIN and node.op_type == OP: |
| if self.stub_scan: |
| |
| stub_node = helper.make_node('Identity', [node.input[0]], [node.output[0]], |
| name=f'depthart_stub_scan_{index}') |
| new_nodes.append(stub_node) |
| index += 1 |
| continue |
| self.node_attrs = node.attribute |
| self.out_name = node.output[0] |
| builder = self.build_logspace if self.scan_method == 'log' else self.build |
| nodes = builder(f'depthart_pscan_{index}', shapes[index], node.input) |
| new_nodes.extend(nodes) |
| index += 1 |
| elif node.op_type == 'Einsum': |
| equation = None |
| for attr in node.attribute: |
| if attr.name == 'equation': |
| equation = attr.s.decode() |
| if equation not in EINSUM_EQUATIONS: |
| raise ValueError(f'{node.name}: 不支持的 Einsum 方程 {equation!r}') |
| x_shape, y_shape = einsum_shapes[node.name] |
| nodes = self.build_einsum(f'depthart_einsum_{einsum_index}', |
| node, x_shape, y_shape) |
| new_nodes.extend(nodes) |
| einsum_index += 1 |
| else: |
| new_nodes.append(node) |
| out = onnx.ModelProto() |
| out.CopyFrom(self.model) |
| del out.graph.node[:] |
| for position, node in enumerate(new_nodes): |
| if not node.name: |
| node.name = f'depthart_gen_{position}' |
| out.graph.node.extend(new_nodes) |
| out.graph.initializer.extend(self.initializers[len(self.model.graph.initializer):]) |
| return out |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser( |
| description='把 DepthART SelectiveScan 改写为 pulsar2 可编译的纯 ONNX 子图') |
| parser.add_argument('--model', type=Path, default=DEFAULT_MODEL, |
| help=f'输入 ONNX，默认 {DEFAULT_MODEL.relative_to(PROJECT_ROOT)}') |
| parser.add_argument('--output', type=Path, |
| default=PROJECT_ROOT / 'onnx' / 'relative_s_448_default_axonnx.onnx', |
| help='输出 ONNX（不会覆盖输入模型）') |
| parser.add_argument('--chunk-size', type=int, default=14, |
| help='分块长度 T，默认 14；越大显存占用越高') |
| parser.add_argument('--stub-scan', action='store_true', |
| help='把 SelectiveScan 替换为 Identity，生成仅用于测 NPU 耗时的对照模型（精度无意义）') |
| parser.add_argument('--scan-method', choices=('pairwise', 'log'), default='pairwise', |
| help='pairwise=两两指数矩阵(默认)；log=对数空间递推，无 T^2 矩阵，T 需小') |
| args = parser.parse_args() |
| logging.basicConfig(level=logging.INFO, format='%(message)s') |
| if args.chunk_size < 1: |
| parser.error('--chunk-size 必须为正数') |
| if not args.model.is_file(): |
| parser.error(f'--model 不存在: {args.model}') |
| if args.output.resolve() == args.model.resolve(): |
| parser.error('--output 不能与 --model 相同') |
|
|
| model = onnx.load(str(args.model)) |
| if not any(node.domain == DOMAIN and node.op_type == OP for node in model.graph.node): |
| raise ValueError(f'{args.model} 中没有 {DOMAIN}::{OP} 节点') |
| rewriter = Rewriter(model, args.chunk_size, args.stub_scan, args.scan_method) |
| rewritten = rewriter.rewrite() |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| onnx.save(rewritten, str(args.output)) |
| logging.info('已写出标准 ONNX: %s（节点数 %d）', args.output, len(rewritten.graph.node)) |
| logging.info('下一步：把 pulsar2 配置的 input 指向该文件；' |
| '可先用 src/infer/depthart_infer_onnx.py 验证数值。') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
