ONNX模型计算图If算子详解

269 阅读1分钟

在ONNX模型计算图中,If是非常特殊的一个算子,在于其必须输入两个子图,然后在运行时根据标量逻辑选择其一运行。ONNX官方文档对If的描述过于简略,以至于很多细节都没有说明白——主要就是子图的输入和输出问题。

顶层图的输入是推理会话在运行时提供的,输出亦由推理会话负责收集。而If的两个子图,同样需要指定输入输出,那么问题就来了,子图的输入如何提供?子图的输出如何收集?文档只说了两个子图和If算子的输出必须数量相同且类型兼容,而没有回答以下重大问题:

子图如何获取输入值?

显然,子图的输入值只能从父图中来。实测发现,子图可以直接引用父图中的名称,不需要任何额外声明。而且,子图的输入声明必须为空,否则运行时将因为未提供输入而失败。

子图的输出如何收集?

不同于输入,子图的输出不是按照名称,而是按照顺序分发的。两个字图的输出都会按照顺序转交给If算子的输出,这个过程中并不要求名称的一致性,只要求输出个数和类型的一致性。