ONNX Script 生成 FunctionProto

以下示例展示了我们如何在 onnxscript 中将 Selu 定义为函数。

首先，导入用于定义函数的 ONNX opset。

from onnxscript import opset15 as op
from onnxscript import script

定义 Selu：

@script()
def Selu(X, alpha: float, gamma: float):
    alphaX = op.CastLike(alpha, X)
    gammaX = op.CastLike(gamma, X)
    neg = gammaX * (alphaX * op.Exp(X) - alphaX)
    pos = gammaX * X
    zero = op.CastLike(0, X)
    return op.Where(zero >= X, neg, pos)

可以像下面这样将 ONNXScript 函数转换为 ONNX 函数（FunctionProto）：

onnx_fun = Selu.to_function_proto()

打印函数：

import onnx  # noqa: E402

print(onnx.printer.to_text(onnx_fun))

<
  domain: "this",
  opset_import: ["" : 15]
>
Selu <alpha,gamma>(X) => (return_val)
{
   alpha = Constant <value_float: float = @alpha> ()
   alphaX = CastLike (alpha, X)
   gamma = Constant <value_float: float = @gamma> ()
   gammaX = CastLike (gamma, X)
   tmp = Exp (X)
   tmp_0 = Mul (alphaX, tmp)
   tmp_1 = Sub (tmp_0, alphaX)
   neg = Mul (gammaX, tmp_1)
   pos = Mul (gammaX, X)
   int64_0 = Constant <value: tensor = int64 int64_0 {0}> ()
   zero = CastLike (int64_0, X)
   tmp_2 = GreaterOrEqual (zero, X)
   return_val = Where (tmp_2, neg, pos)
}