onnxscript_linear_model_definition_and_onnxruntime_execution.py

1import numpy as np
2import onnxruntime as ort
3from onnxscript import opset15 as op
4from onnxscript import script
5
6# 1. Define the ONNX model using Python syntax
7@script()
8def linear_model(A: float, X: float, B: float) -> float:
9    return op.MatMul(A, X) + B
10
11# 2. Convert the script to an ONNX ModelProto
12onnx_model = linear_model.to_model_proto()
13
14# 3. Prepare data for testing
15a = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
16x = np.array([[5.0, 6.0], [7.0, 8.0]], dtype=np.float32)
17b = np.array([[1.0, 1.0], [1.0, 1.0]], dtype=np.float32)
18
19# 4. Run the model using ONNX Runtime
20sess = ort.InferenceSession(onnx_model.SerializeToString(), providers=["CPUExecutionProvider"])
21inputs = {
22    "A": a,
23    "X": x,
24    "B": b
25}
26outputs = sess.run(None, inputs)
27
28print(outputs[0])