spox实现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