生成 ModelProto

此示例展示了如何使用 onnxscript 定义 ONNX 模型。onnxscript 的行为类似于编译器。它将 script 转换为 ONNX 模型。

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

import numpy as np
import onnx
from onnxruntime import InferenceSession

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


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

可以按照以下方式将其转换为模型（ ONNX ModelProto）：

model = square_loss.to_model_proto()

打印：

print(onnx.printer.to_text(model))

<
   ir_version: 8,
   opset_import: ["" : 15]
>
square_loss (float[N,1] X, float[N,1] Y) => (float[1,1] return_val) {
   diff = Sub (X, Y)
   tmp = Mul (diff, diff)
   return_val = ReduceSum <keepdims: int = 1> (tmp)
}

可以使用标准的 ONNX API 对模型进行形状推断和类型检查。

model = onnx.shape_inference.infer_shapes(model)
onnx.checker.check_model(model)

最后，我们可以使用标准的 onnxruntime API 通过 onnxruntime 计算该模型的输出。

sess = InferenceSession(model.SerializeToString(), providers=("CPUExecutionProvider",))

X = np.array([[0, 1, 2]], dtype=np.float32).T
Y = np.array([[0.1, 1.2, 2.3]], dtype=np.float32).T

got = sess.run(None, {"X": X, "Y": Y})
expected = ((X - Y) ** 2).sum()

print(expected, got)

0.13999999 [array([[0.13999999]], dtype=float32)]