ONNX ModelProto 属性

ONNX ModelProto 属性#

在 ONNX 中,ModelProto 通常存储计算图之外的额外信息,例如 ir_versionproducer_name。通过向 script(或 to_model_proto)调用传递额外的命名参数,可以设置生成的 ModelProto 的这些属性,如下例所示。只有以这种方式指定的 protobuf 消息 ModelProto 中定义的有效字段才应被指定。

首先,我们在 onnxscript 中定义一个平方损失函数的实现。

import onnx

from onnxscript import FLOAT, script
from onnxscript import opset15 as op


@script(ir_version=7, producer_name="OnnxScript", producer_version="0.1")
def square_loss(X: FLOAT["N"], Y: FLOAT["N"]) -> FLOAT[1]:  # noqa: F821
    diff = X - Y
    return op.ReduceSum(diff * diff, keepdims=1)

打印:

model = square_loss.to_model_proto()
print(onnx.printer.to_text(model))
<
   ir_version: 7,
   opset_import: ["" : 15],
   producer_name: "OnnxScript",
   producer_version: "0.1"
>
square_loss (float[N] X, float[N] Y) => (float[1] return_val) {
   diff = Sub (X, Y)
   tmp = Mul (diff, diff)
   return_val = ReduceSum <keepdims: int = 1> (tmp)
}