onnxscript_relu_function_definition_and_model_export.py

1import onnxscript
2from onnxscript import opset15 as op
3
4# Define an ONNX function using the @script decorator
5@onnxscript.script()
6def relu(x):
7    # Use standard ONNX operators from the opset
8    return op.Where(x > 0, x, 0.0)
9
10# Convert the function to an ONNX ModelProto
11model_proto = relu.to_model_proto()
12
13# Save the model to a file
14import onnx
15onnx.save(model_proto, "relu.onnx")
16
17# Example of inspecting the generated ONNX code
18print(onnx.helper.printable_graph(model_proto.graph))