spox是直接创建onnx模型,或多个onnx模型整合为1个模型的工具。
由于pytorch导出的onnx没有if或loop操作。
本例主要讲spox实现loop循环。
代码
简单的例子
这里实现一个0+1+2+3+…+100的样例
(sum_x,) = op.loop(
M=const(101), # 注意是101,
v_initial=[const(0)], # a := 0 at the start
body=lambda i, _, a: [ # iteration (i), 本轮的条件状态 (_), accumulator (a)
const(True), # continue
op.add(a, i) # step is a := a+i
]
)
M是指迭代次数,iterations是不可控的自增,从0开始,所以最后的iteration值为n-1。
M也可以不配置,由body返回的第一个条件控制是否继续循环。
更全面的了解各个参数
It has 2+N inputs: (iteration_num, condition, loop carried dependencies...). It has 1+N+K outputs: (condition, loop carried dependencies..., scan_outputs...)
主要了解迭代参数与
def body_(iteration, _, a1, a2): # iteration轮次,从0开始的index。 a1 a2是本轮输入。 第2个参数为本轮状态始终为true
s1 = iteration#op.arg_max(x, axis=-1) # 每轮独立记录的值
s2 = const(1) # 每轮独立记录的值
condition = const(True) # 提前终止的条件
a1 = op.add(a1, const(1)) # 循环给下轮的值
a2 = op.sub(a2, const(1)) # 循环给下轮的值
return condition, a1, a2, s1, s2
# a1, a2, a3 可以为任意数量, s1, s2可以为任意数量也可以没有
(a1, a2, s1, s2) = op.loop(
M=const(4), # x times iterations 遍例多少次,可以为空,让body里的条件来控制
v_initial=[const([2, 3, 4]), const([2, 3])], # a1 a2的初始值
body=body_
)
输出:
dtype: int64 , shape: (3,) , value: [6 7 8]
dtype: int64 , shape: (2,) , value: [-2 -1]
dtype: int64 , shape: (4, 1) , value: [[0]
[1]
[2]
[3]]
dtype: int64 , shape: (4,) , value: [1 1 1 1]
完整代码如下:
import warnings
import logging
import numpy as np
import onnx
import onnxruntime
import spox._future
from spox import argument, build, Tensor, Var, inline
import spox.opset.ai.onnx.v17 as op # op - ai.onnx@17
def const(value):
return op.constant(value=np.array(value))
def scalar(var: Var):
return op.reshape(var, const(np.array([], int)))
def run(model: onnx.ModelProto, **kwargs) -> list[np.ndarray]:
options = onnxruntime.SessionOptions()
options.log_severity_level = 3
return onnxruntime.InferenceSession(model.SerializeToString(), options).run(
None,
{k: np.array(v) for k, v in kwargs.items()}
)
warnings.filterwarnings("ignore")
logging.basicConfig(level=logging.DEBUG)
spox._future.set_value_prop_backend(spox._future.ValuePropBackend.ONNXRUNTIME)
x = argument(Tensor(float, ('N', 'N')))
# It has 2+N inputs: (iteration_num, condition, loop carried dependencies...). It has 1+N+K outputs: (condition, loop carried dependencies..., scan_outputs...)
def b_(iteration, _, a1, a2): # iteration轮次,从0开始的index。 a1 a2是本轮输入。 第2个参数为本轮状态始终为true
s1 = iteration#op.arg_max(x, axis=-1) # 每轮独立记录的值
s2 = const(1) # 每轮独立记录的值
condition = const(True) # 提前终止的条件
a1 = op.add(a1, const(1)) # 循环给下轮的值
a2 = op.sub(a2, const(1)) # 循环给下轮的值
return condition, a1, a2, s1, s2
# a1, a2, a3 可以为任意数量, s1, s2可以为任意数量也可以没有
(a1, a2, s1, s2) = op.loop(
M=const(4), # x times iterations 遍例多少次,可以为空,让body里的条件来控制
v_initial=[const([2, 3, 4]), const([2, 3])], # a1 a2的初始值
body=b_
)
fact_model = build({"x": x}, {"a1": a1, "a2": a2, "s1": s1, 's2':s2})
rr = run(fact_model, x=np.array([[1., 2., 3., 4.], [9., 8., 7., 6.], [19., 18., 17., 16.]]))
for i in rr:
print('dtype:', i.dtype, ', shape:', i.shape, ', value:', i)
a1, a2, a3 可以为任意数量, s1, s2可以为任意数量也可以没有