ONNX ModelProto 属性#
在 ONNX 中,ModelProto 通常存储计算图之外的额外信息,例如 ir_version
或 producer_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)
}