What, Why and How ONNX Script?
ONNX (Open Neural Network eXchange) is an open-source standard that defines a common set of operators and a file format that enables AI developers to use models with various frameworks, tools, runtimes, and compilers. For a deeper dive into ONNX check this post!
In this article, I am excited to share some thoughts on @Microsoft’s recent announcement on ONNX Script and what it enables developers to do.
ONNX Script enables developers to express ONNX functions and models in Python. This simple yet powerful expression of models/ functions allows for eager mode evaluation thus significantly improving debug ability.
ONNX Script provides developers with the following capabilities:
- A converter to translate a Python ONNX Script function into an ONNX graph.
- A converter that translates ONNX models and functions into ONNX Script. This capability can be used to fully round-trip.
- A runtime(ONNXRuntime) that allows eager mode execution of the model.
For those familiar with the torch ecosystem, TorchScript provides a way to create serializable and optimizable models from PyTorch code that can be saved from a Python process and loaded in a process where there is no Python dependency. This makes it possible to train models in PyTorch using familiar tools in Python and then export the model via TorchScript to a production environment where Python programs may be disadvantageous for performance and multi-threading reasons.
ONNX Script essentially provides the Python-based front end to the existing ONNX infrastructure and also the tools to translate the script to serialized ONNX graphs that can be deployed without dependencies.
Why should you pay attention to ONNX Script?
ONNX models were traditionally obtained by exporting models built in different source frameworks such as PyTorch and TensorFlow using conversion tools provided by respective frameworks. This has allowed developers to stick to their existing development flow in PyTorch/ TensorFlow and then convert the models to ONNX for deployment to take advantage of the flexibility ONNX provides.
But the downside to this approach is the issues faced during model conversion from the source framework to ONNX. This is due to a mismatch in operator coverage surface area between the source frameworks(PyTorch/ TensorFlow) and ONNX.
To solve this issue, instead of brute forcing a 1:1 operator match with the source frameworks, ONNX approaches this problem by breaking down unsupported complex operations into their fundamental components and expressing them in the decomposed form of simpler existing operators through ONNX functions.
ONNX Script provides a new way to build, understand and debug complex ONNX functions/models in a way that developers are already used to, which will reduce the inertia in adoption.
This enables Domain-Specific Hardware accelerators and compiler developers like Groq to focus on a tight set of operations and represent the broader ecosystem of models using this optimized set of fundamental operations.
How can you use ONNX-script?
The following are a few ways you can start using ONNX script.
- Represent complex operations using existing ONNX primitives
The following example shows how a Gaussian Error Linear Units (GELU) activation can be represented using ONNX Script.
The GELU nonlinearity weights inputs by their percentile, rather than gates inputs by their sign as in ReLUs. GELU is mathematically expressed as follows:
Using the above mathematical expression GELU can be represented in ONNX Script as follows
Explore more examples here!
2. Author ONNX models that can be exported for training/inference.
ONNX Script enables developers to write entire models in Python and export them into an ONNX model using the following:
from onnxscript import script
from onnxscript.onnx_types import FLOAT
import onnx@script()
def model(input: FLOAT[...]):
....
....
....model_proto = model.to_model_proto()
onnx.save(model_proto, "path/model.onnx")
Note: Currently a function with attributes cannot be exported as a model i.e. the function cannot take arbitrary variables that can be used inside the function.
3. Potentially replace operations in source frameworks that block conversion to ONNX
For example, currently, PyTorch models that use torch.bucketize like SuperGATConv from PyTorch Geometric are blocked from being converted to ONNX due to the lack of an equivalent operator on the ONNX. ONNX Script can be used to potentially resolve this as follows:
# pytorch-model.pyimport torch
import torch.nn as nn
from onnxscript import script
from onnxscript.onnx_types import FLOATclass model(nn.Module):
def __init__(self):
super(model, self).__init__() def forward(self):
....
....
torch.bucketize(...)
....@script()
def onnx_bucketize(input: FLOAT[...]):
....
....
....# Replace torch.bucketize(...) with onnx_bucketize(...)
torch_model = model(...)
torch.onnx.export(torch_model, inputs, "model.onnx")
ONNX Script has the potential to significantly improve the developer experience in the ONNX ecosystem. Though this is still the early days of ONNX Script, I am personally excited to see bridges this builds/ strengthens in this fragmented world of AI frameworks!!
