ONNX Script 急切模式评估

onnxscript 函数可以直接作为 Python 函数执行（例如，使用 Python 调试器）。这对于调试 onnxscript 函数定义非常有用。这种执行利用了函数定义中使用的 ONNX 操作的后端实现。目前，后端实现使用 onnxruntime 来执行每个操作调用。这种执行模式被称为急切模式评估。

下面的示例说明了这一点。我们首先定义一个 onnxscript 函数：

import numpy as np

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


@script()
def linear(A: FLOAT["N", "K"], W: FLOAT["K", "M"], Bias: FLOAT["M"]) -> FLOAT["N", "M"]:  # noqa: F821
    T1 = op.MatMul(A, W)
    T2 = op.Add(T1, Bias)
    Y = op.Relu(T2)
    return Y

创建用于评估函数的输入：

np.random.seed(0)
m = 4
k = 16
n = 4
a = np.random.rand(k, m).astype("float32").T
w = np.random.rand(n, k).astype("float32").T
b = np.random.rand(n).astype("float32").T

评估结果：

print(linear(a, w, b))

[[3.7695956 3.8361263 5.116064  5.2047744]
 [4.5182567 4.2103305 4.54666   5.6752048]
 [4.0728097 3.1566992 4.821034  4.7809625]
 [4.925565  3.558134  4.7679787 5.3899584]]