1TorchDynamo-based ONNX Exporter 2=============================== 3 4.. automodule:: torch.onnx 5 :noindex: 6 7.. contents:: :local: 8 :depth: 3 9 10.. warning:: 11 The ONNX exporter for TorchDynamo is a rapidly evolving beta technology. 12 13Overview 14-------- 15 16The ONNX exporter leverages TorchDynamo engine to hook into Python's frame evaluation API 17and dynamically rewrite its bytecode into an FX Graph. 18The resulting FX Graph is then polished before it is finally translated into an ONNX graph. 19 20The main advantage of this approach is that the `FX graph <https://pytorch.org/docs/stable/fx.html>`_ is captured using 21bytecode analysis that preserves the dynamic nature of the model instead of using traditional static tracing techniques. 22 23The exporter is designed to be modular and extensible. It is composed of the following components: 24 25 - **ONNX Exporter**: :class:`Exporter` main class that orchestrates the export process. 26 - **ONNX Export Options**: :class:`ExportOptions` has a set of options that control the export process. 27 - **ONNX Registry**: :class:`OnnxRegistry` is the registry of ONNX operators and functions. 28 - **FX Graph Extractor**: :class:`FXGraphExtractor` extracts the FX graph from the PyTorch model. 29 - **Fake Mode**: :class:`ONNXFakeContext` is a context manager that enables fake mode for large scale models. 30 - **ONNX Program**: :class:`ONNXProgram` is the output of the exporter that contains the exported ONNX graph and diagnostics. 31 - **ONNX Diagnostic Options**: :class:`DiagnosticOptions` has a set of options that control the diagnostics emitted by the exporter. 32 33Dependencies 34------------ 35 36The ONNX exporter depends on extra Python packages: 37 38 - `ONNX <https://onnx.ai>`_ 39 - `ONNX Script <https://onnxscript.ai>`_ 40 41They can be installed through `pip <https://pypi.org/project/pip/>`_: 42 43.. code-block:: bash 44 45 pip install --upgrade onnx onnxscript 46 47`onnxruntime <https://onnxruntime.ai>`_ can then be used to execute the model 48on a large variety of processors. 49 50A simple example 51---------------- 52 53See below a demonstration of exporter API in action with a simple Multilayer Perceptron (MLP) as example: 54 55.. code-block:: python 56 57 import torch 58 import torch.nn as nn 59 60 class MLPModel(nn.Module): 61 def __init__(self): 62 super().__init__() 63 self.fc0 = nn.Linear(8, 8, bias=True) 64 self.fc1 = nn.Linear(8, 4, bias=True) 65 self.fc2 = nn.Linear(4, 2, bias=True) 66 self.fc3 = nn.Linear(2, 2, bias=True) 67 68 def forward(self, tensor_x: torch.Tensor): 69 tensor_x = self.fc0(tensor_x) 70 tensor_x = torch.sigmoid(tensor_x) 71 tensor_x = self.fc1(tensor_x) 72 tensor_x = torch.sigmoid(tensor_x) 73 tensor_x = self.fc2(tensor_x) 74 tensor_x = torch.sigmoid(tensor_x) 75 output = self.fc3(tensor_x) 76 return output 77 78 model = MLPModel() 79 tensor_x = torch.rand((97, 8), dtype=torch.float32) 80 onnx_program = torch.onnx.export(model, (tensor_x,), dynamo=True) 81 82As the code above shows, all you need is to provide :func:`torch.onnx.export` with an instance of the model and its input. 83The exporter will then return an instance of :class:`torch.onnx.ONNXProgram` that contains the exported ONNX graph along with extra information. 84 85The in-memory model available through ``onnx_program.model_proto`` is an ``onnx.ModelProto`` object in compliance with the `ONNX IR spec <https://github.com/onnx/onnx/blob/main/docs/IR.md>`_. 86The ONNX model may then be serialized into a `Protobuf file <https://protobuf.dev/>`_ using the :meth:`torch.onnx.ONNXProgram.save` API. 87 88.. code-block:: python 89 90 onnx_program.save("mlp.onnx") 91 92Two functions exist to export the model to ONNX based on TorchDynamo engine. 93They slightly differ in the way they produce the :class:`ExportedProgram`. 94:func:`torch.onnx.dynamo_export` was introduced with PyTorch 2.1 and 95:func:`torch.onnx.export` was extended with PyTorch 2.5 to easily switch 96from TorchScript to TorchDynamo. To call the former function, 97the last line of the previous example can be replaced by the following one. 98 99.. code-block:: python 100 101 onnx_program = torch.onnx.dynamo_export(model, tensor_x) 102 103Inspecting the ONNX model using GUI 104----------------------------------- 105 106You can view the exported model using `Netron <https://netron.app/>`__. 107 108.. image:: _static/img/onnx/onnx_dynamo_mlp_model.png 109 :width: 40% 110 :alt: MLP model as viewed using Netron 111 112Note that each layer is represented in a rectangular box with a *f* icon in the top right corner. 113 114.. image:: _static/img/onnx/onnx_dynamo_mlp_model_function_highlight.png 115 :width: 40% 116 :alt: ONNX function highlighted on MLP model 117 118By expanding it, the function body is shown. 119 120.. image:: _static/img/onnx/onnx_dynamo_mlp_model_function_body.png 121 :width: 50% 122 :alt: ONNX function body 123 124The function body is a sequence of ONNX operators or other functions. 125 126When the conversion fails 127------------------------- 128 129Function :func:`torch.onnx.export` should called a second time with 130parameter ``report=True``. A markdown report is generated to help the user 131to resolve the issue. 132 133Function :func:`torch.onnx.dynamo_export` generates a report using 'SARIF' format. 134ONNX diagnostics goes beyond regular logs through the adoption of 135`Static Analysis Results Interchange Format (aka SARIF) <https://docs.oasis-open.org/sarif/sarif/v2.1.0/sarif-v2.1.0.html>`__ 136to help users debug and improve their model using a GUI, such as 137Visual Studio Code's `SARIF Viewer <https://marketplace.visualstudio.com/items?itemName=MS-SarifVSCode.sarif-viewer>`_. 138 139The main advantages are: 140 141 - The diagnostics are emitted in machine parseable `Static Analysis Results Interchange Format (SARIF) <https://docs.oasis-open.org/sarif/sarif/v2.1.0/sarif-v2.1.0.html>`__. 142 - A new clearer, structured way to add new and keep track of diagnostic rules. 143 - Serve as foundation for more future improvements consuming the diagnostics. 144 145.. toctree:: 146 :maxdepth: 1 147 :caption: ONNX Diagnostic SARIF Rules 148 :glob: 149 150 generated/onnx_dynamo_diagnostics_rules/* 151 152API Reference 153------------- 154 155.. autofunction:: torch.onnx.dynamo_export 156 157.. autoclass:: torch.onnx.ExportOptions 158 :members: 159 160.. autofunction:: torch.onnx.enable_fake_mode 161 162.. autoclass:: torch.onnx.ONNXProgram 163 :members: 164 165.. autoclass:: torch.onnx.ONNXRuntimeOptions 166 :members: 167 168.. autoclass:: torch.onnx.OnnxExporterError 169 :members: 170 171.. autoclass:: torch.onnx.OnnxRegistry 172 :members: 173 174.. autoclass:: torch.onnx.DiagnosticOptions 175 :members: 176
