Typically, machine learning models are developed using higher-level frameworks such as\u202fPyTorch<\/a>\u202fand\u202fTensorFlow<\/a>. While these frameworks tend to be productive for iterating on the development of models, the models are not typically deployed to production in this fashion. Instead, they are exported to ONNX by facilities provided by the frameworks, and then optimized for a particular target by tools such as\u202fOlive<\/a>. <\/p>\n\n\n\n Want to dive right into ONNX Script? Check it out on GitHub<\/a>.<\/p>\n\n\n\n ONNX Script is a new open-source library for directly authoring ONNX models in Python with a focus on\u202fclean, idiomatic Python\u202fsyntax and composability through ONNX-native functions. Critically, it is also the foundation upon which we are building the new PyTorch ONNX exporter to support\u202fTorchDynamo<\/a>\u2014the future of PyTorch. <\/p>\n\n\n\n Prior to ONNX Script, authoring ONNX models required deep knowledge of the specification and serialization format itself. While eventually a more convenient\u202fhelper API<\/a>\u202fwas introduced that largely abstracted the serialization format, it still required deep familiarity with ONNX constructs. <\/p>\n\n\n\n ONNX Script takes a new approach by integrating deeply with Python on two fronts:<\/p>\n\n\n\n Let\u2019s look at how we might implement GELU using ONNX Script and compare it with\u202fonnx.helper\u202fAPI\u2014and to be clear\u2014the examples below produce the same ONNX model. For reference, we\u2019ll use this definition of GELU to guide the ONNX implementations: <\/p>\n\n\n As you can see, ONNX Script emphasizes the familiar readability and productivity of Python while expressing an ONNX model that can be statically reasoned about by existing Python and ONNX tooling. <\/p>\n\n\n\n This also means ONNX comes alive within the context of your existing tooling and development environments, be it debugging in Visual Studio Code or demonstrating concepts in a Jupyter Notebook\u2014ONNX Script integrates naturally. <\/p>\n\n\n\n Much has changed since ONNX support for PyTorch was originally introduced over five years ago in PyTorch 0.3.0. For PyTorch 2.0, TorchDynamo represents the eventual deprecation of TorchScript, which implies a major overhaul of the ONNX exporter is necessary. We are fully embracing this as an opportunity to revisit the fundamentals upon which the exporter is built, and ONNX Script is its new foundation. We began this effort in November 2022 and have worked closely with PyTorch engineers to ensure TorchDynamo is a fully capable starting point for exporting high-fidelity ONNX for years to come. <\/p>\n\n\n\n One of the first streams of work we started was the development of what we call Torchlib, a pure ONNX implementation of the operators in PyTorch\u2014namely\u202fCore ATen IR and Prims IR<\/a>\u2014and of course, these operators are implemented in ONNX Script. This approach greatly simplifies the central responsibility of the exporter as it \u201cjust\u201d needs to project FX graph nodes produced by TorchDynamo into ONNX graph nodes, without concerning itself with the implementation details of individual operators. <\/p>\n\n\n\n We will cover the new PyTorch ONNX exporter in a separate post with more depth as PyTorch 2.1 approaches, but for now, the key takeaway is that ONNX Script is pervasive throughout our renewed approach. For those willing to try bold new things, the new exporter is available as a preview in PyTorch nightly via the\u202ftorch.onnx.dynamo_export<\/a>\u202fAPI. <\/p>\n\n\n\n By deeply weaving ONNX Script support into the PyTorch ONNX exporter, we have also made it possible to augment PyTorch model code with specialized ONNX functions as custom operators. We introduced initial support for this in the TorchScript exporter starting with\u202fPyTorch 1.13<\/a>\u202fand continue to refine this capability in the new exporter. <\/p>\n\n\n\n Let\u2019s look at an example slightly more complicated than GELU. In fact, the following example is adapted directly from the new PyTorch ONNX exporter, implementing support for\u202ftorch.chunk<\/a>, which attempts to split a tensor into the number of specified chunks. <\/p>\n\n\n We start by importing from onnxscript the ONNX opset we wish to use (version 18 in this case), the\u202f Next, we define the\u202f We also need to define a simple model that calls our ONNX Script function so we can export and verify an end-to-end example: <\/p>\n\n\n This model will simply split the provided tensor into ten tensors, but it also demonstrates that ONNX functions can of course call other ONNX functions, not just built-in ONNX operators. <\/p>\n\n\n\n We\u2019ll now export our ONNX Script model to ONNX and explore it in Netron. Functions decorated with If we open\u202ften_chunks_model.onnx<\/em>\u202fin Netron, we can observe the composability of ONNX functions and how the Python code was translated into an ONNX model. <\/p>\n\n\n The graphs depict our two ONNX functions. We can observe the original\u202finput tensor<\/strong>\u202fflowing from\u202f Finally, we should test our model. ONNX Script makes this easy since it provides a mechanism for eagerly evaluating the model through either\u202fONNX Runtime<\/a>\u202for the new\u202fONNX Reference Evaluator<\/a>. Of note, ONNX Script has built-in support for\u202fNumPy<\/a>\u202ffor input and output values to ease the overhead of creating and filling tensors in ONNX. <\/p>\n\n\n Because ONNX Script\u2019s eager mode evaluates the model on an op-by-op basis, it is conducive to debugging ONNX using standard Python tooling such as\u202fpdb<\/a>\u202fdirectly or through richer IDE and editor integrations provided by Visual Studio and Visual Studio Code. <\/p>\n\n\n\n Along with debuggability, IntelliSense support is front-and-center with ONNX Script. We rely heavily on Python type annotations to enforce the correctness of ONNX and to make ONNX more discoverable, including inline documentation on hover tooltips and code completion suggestions. A single click will take you to expanded online documentation in your browser as well. <\/p>\n\n\n\n In summary, ONNX Script offers a new Python-first programming model for authoring ONNX models that integrates with the existing rich ecosystem of Python tooling and environments. <\/p>\n\n\n\n Going forward, we envision ONNX Script as a means for defining and extending ONNX itself. New core operators and higher-order functions that are intended to become part of the ONNX standard could be authored in ONNX Script as well, reducing the time and effort it takes for the standard to evolve. We have proven this is a viable strategy by developing Torchlib, upon which the upcoming PyTorch Dynamo-based ONNX exporter is built. <\/p>\n\n\n\n Over the coming months, we will also support converting ONNX into ONNX Script to enable seamless editing of existing models, which can play a key role in optimization passes, but also allow for maintaining and evolving ONNX models more naturally. We also intend to propose ONNX Script for inclusion directly within the ONNX GitHub organization soon, under the Linux Foundation umbrella. <\/p>\n\n\n\n Check out ONNX Script today on GitHub<\/a>!<\/p>\n\n\n\nAnnouncing ONNX Script <\/h2>\n\n\n\n
\n
a\u202f+\u202fb<\/code>\u202fin Python would translate to the ONNX operator\u202fas Add(a,\u202fb)<\/code>.<\/li>\n<\/ol>\n\n\n\n![\"GELU](\"https:\/\/opensource.microsoft.com\/blog\/wp-content\/uploads\/2023\/07\/gelu_onnxscript_onnxhelper.webp\")
<\/figure>\n\n\n\nWhy are we investing in ONNX Script? <\/h2>\n\n\n\n
An end-to-end example <\/h2>\n\n\n\n
![\"The](\"https:\/\/opensource.microsoft.com\/blog\/wp-content\/uploads\/2023\/07\/aten_chunk.webp\")
<\/figure>\n\n\n\n@script<\/code>\u202fdecorator, and the tensor types of\u202fFLOAT<\/code>\u202fand\u202fINT64<\/code>. In ONNX Script, tensor shapes are denoted by subscripting the type, such as FLOAT[2, 10]<\/code>, symbolically such as FLOAT[\"M\", \"N\"]<\/code>, or FLOAT[...]<\/code> incase the tensor shape is unknown. Without subscripting (just FLOAT<\/code>), the type is intended to indicate a scalar (a tensor of rank 0).<\/p>\n\n\n\naten_chunk<\/code>\u202ffunction with type annotations and implement the body of the function using both built-in Python syntax and explicit invocations of ONNX operators. The example uses various binary expressions and an\u202fif\u202fstatement, but many other idiomatic Python constructs are also supported. <\/p>\n\n\n\n![\"ten_chunks_model()](\"https:\/\/opensource.microsoft.com\/blog\/wp-content\/uploads\/2023\/07\/ten_chunks_model.webp\")
<\/figure>\n\n\n\n@script<\/code> allow them to be exported using the to_model_proto<\/code> function.<\/p>\n\n\n![\"saving](\"https:\/\/opensource.microsoft.com\/blog\/wp-content\/uploads\/2023\/07\/ten_chunks_model_to_onnx.webp\")
<\/figure>\n\n\n\n![\"a](\"https:\/\/opensource.microsoft.com\/blog\/wp-content\/uploads\/2023\/06\/Picture1-1.webp\")
<\/figure>\n\n\n\nten_chunks_model<\/code>\u202finto\u202faten_chunk<\/code>\u202falong with the\u202fattribute chunks=10<\/strong>. A\u202fsequence of \u2264 10 tensors<\/strong>\u202fis returned. As one might expect, functions in ONNX can be defined once and invoked any number of times within a model.\u202fRead more about core ONNX concepts<\/a>.<\/em> <\/p>\n\n\n\nIterating and debugging <\/h2>\n\n\n\n
![\"setting](\"https:\/\/opensource.microsoft.com\/blog\/wp-content\/uploads\/2023\/07\/ten_chunks_model_eager_eval.webp\")
<\/figure>\n\n\n![\"Output](\"https:\/\/opensource.microsoft.com\/blog\/wp-content\/uploads\/2023\/07\/ten_chunks_model_eager_eval_output.webp\")
<\/figure>\n\n\n\n![\"A](\"https:\/\/opensource.microsoft.com\/blog\/wp-content\/uploads\/2023\/06\/Picture2-1.png\")
dim_size<\/code>\u202fvariable and call stack.<\/em><\/figcaption><\/figure>\n\n\n\n![\"A](\"https:\/\/opensource.microsoft.com\/blog\/wp-content\/uploads\/2023\/06\/Picture3-2.png\")
What\u2019s next for ONNX Script? <\/h2>\n\n\n\n
\n\n\n\n