spox实现onnx模型中的for-loop循环

spox

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可以为任意数量也可以没有