生成 LibProto

生成 LibProto#

以下示例展示了我们如何定义一个由多个函数组成的库,并将其导出。

这是原型的。要完全支持 LibProto,需要原型扩展。

from onnxscript import export_onnx_lib, script
from onnxscript import opset15 as op
from onnxscript.values import Opset

下面定义的库函数的域/版本

opset = Opset("com.mydomain", 1)

定义函数库:

@script(opset)
def l2norm(X):
    return op.ReduceSum(X * X, keepdims=1)


@script(opset)
def square_loss(X, Y):
    return l2norm(op.Sub(X, Y))

将函数作为 ONNX 库导出。

export_onnx_lib([l2norm, square_loss], "mylib.onnxlib")