onnx-array-api implements a numpy API for ONNX. It gives the user the ability to convert functions written following the numpy API to convert that function into ONNX as well as to execute it.
import numpy as np
from onnx_array_api.npx import absolute, jit_onnx
from onnx_array_api.plotting.text_plot import onnx_simple_text_plot
def l1_loss(x, y):
return absolute(x - y).sum()
def l2_loss(x, y):
return ((x - y) ** 2).sum()
def myloss(x, y):
return l1_loss(x[:, 0], y[:, 0]) + l2_loss(x[:, 1], y[:, 1])
jitted_myloss = jit_onnx(myloss)
x = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
y = np.array([[0.11, 0.22], [0.33, 0.44]], dtype=np.float32)
res = jitted_myloss(x, y)
print(res)
print(onnx_simple_text_plot(jitted_myloss.get_onnx()))[0.042]
opset: domain='' version=18
input: name='x0' type=dtype('float32') shape=['', '']
input: name='x1' type=dtype('float32') shape=['', '']
Sub(x0, x1) -> r__0
Abs(r__0) -> r__1
ReduceSum(r__1, keepdims=0) -> r__2
output: name='r__2' type=dtype('float32') shape=None
It supports eager mode as well:
import numpy as np
from onnx_array_api.npx import absolute, eager_onnx
def l1_loss(x, y):
err = absolute(x - y).sum()
print(f"l1_loss={err.numpy()}")
return err
def l2_loss(x, y):
err = ((x - y) ** 2).sum()
print(f"l2_loss={err.numpy()}")
return err
def myloss(x, y):
return l1_loss(x[:, 0], y[:, 0]) + l2_loss(x[:, 1], y[:, 1])
eager_myloss = eager_onnx(myloss)
x = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
y = np.array([[0.11, 0.22], [0.33, 0.44]], dtype=np.float32)
res = eager_myloss(x, y)
print(res)l1_loss=[0.04] l2_loss=[0.002] [0.042]
The library is released on pypi/onnx-array-api and its documentation is published at (Numpy) Array API for ONNX.