• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1torch.onnx
2==========
3
4Overview
5--------
6
7`Open Neural Network eXchange (ONNX) <https://onnx.ai/>`_ is an open standard
8format for representing machine learning models. The ``torch.onnx`` module captures the computation graph from a
9native PyTorch :class:`torch.nn.Module` model and converts it into an
10`ONNX graph <https://github.com/onnx/onnx/blob/main/docs/IR.md>`_.
11
12The exported model can be consumed by any of the many
13`runtimes that support ONNX <https://onnx.ai/supported-tools.html#deployModel>`_, including
14Microsoft's `ONNX Runtime <https://www.onnxruntime.ai>`_.
15
16**There are two flavors of ONNX exporter API that you can use, as listed below.**
17Both can be called through function :func:`torch.onnx.export`.
18Next example shows how to export a simple model.
19
20.. code-block:: python
21
22    import torch
23
24    class MyModel(torch.nn.Module):
25        def __init__(self):
26            super(MyModel, self).__init__()
27            self.conv1 = torch.nn.Conv2d(1, 128, 5)
28
29        def forward(self, x):
30            return torch.relu(self.conv1(x))
31
32    input_tensor = torch.rand((1, 1, 128, 128), dtype=torch.float32)
33
34    model = MyModel()
35
36    torch.onnx.export(
37        model,                  # model to export
38        (input_tensor,),        # inputs of the model,
39        "my_model.onnx",        # filename of the ONNX model
40        input_names=["input"],  # Rename inputs for the ONNX model
41        dynamo=True             # True or False to select the exporter to use
42    )
43
44Next sections introduces the two versions of the exporter.
45
46TorchDynamo-based ONNX Exporter
47-------------------------------
48
49*The TorchDynamo-based ONNX exporter is the newest (and Beta) exporter for PyTorch 2.1 and newer*
50
51TorchDynamo engine is leveraged to hook into Python's frame evaluation API and dynamically rewrite its
52bytecode into an FX Graph. The resulting FX Graph is then polished before it is finally translated into an
53ONNX graph.
54
55The main advantage of this approach is that the `FX graph <https://pytorch.org/docs/stable/fx.html>`_ is captured using
56bytecode analysis that preserves the dynamic nature of the model instead of using traditional static tracing techniques.
57
58:doc:`Learn more about the TorchDynamo-based ONNX Exporter <onnx_dynamo>`
59
60TorchScript-based ONNX Exporter
61-------------------------------
62
63*The TorchScript-based ONNX exporter is available since PyTorch 1.2.0*
64
65`TorchScript <https://pytorch.org/docs/stable/jit.html>`_ is leveraged to trace (through :func:`torch.jit.trace`)
66the model and capture a static computation graph.
67
68As a consequence, the resulting graph has a couple limitations:
69
70* It does not record any control-flow, like if-statements or loops;
71* Does not handle nuances between ``training`` and ``eval`` mode;
72* Does not truly handle dynamic inputs
73
74As an attempt to support the static tracing limitations, the exporter also supports TorchScript scripting
75(through :func:`torch.jit.script`), which adds support for data-dependent control-flow, for example. However, TorchScript
76itself is a subset of the Python language, so not all features in Python are supported, such as in-place operations.
77
78:doc:`Learn more about the TorchScript-based ONNX Exporter <onnx_torchscript>`
79
80Contributing / Developing
81-------------------------
82
83The ONNX exporter is a community project and we welcome contributions. We follow the
84`PyTorch guidelines for contributions <https://github.com/pytorch/pytorch/blob/main/CONTRIBUTING.md>`_, but you might
85also be interested in reading our `development wiki <https://github.com/pytorch/pytorch/wiki/PyTorch-ONNX-exporter>`_.
86
87.. toctree::
88    :hidden:
89
90    onnx_dynamo
91    onnx_dynamo_onnxruntime_backend
92    onnx_torchscript
93
94.. This module needs to be documented. Adding here in the meantime
95.. for tracking purposes
96.. py:module:: torch.onnx.errors
97.. py:module:: torch.onnx.operators
98.. py:module:: torch.onnx.symbolic_caffe2
99.. py:module:: torch.onnx.symbolic_helper
100.. py:module:: torch.onnx.symbolic_opset10
101.. py:module:: torch.onnx.symbolic_opset11
102.. py:module:: torch.onnx.symbolic_opset13
103.. py:module:: torch.onnx.symbolic_opset14
104.. py:module:: torch.onnx.symbolic_opset15
105.. py:module:: torch.onnx.symbolic_opset16
106.. py:module:: torch.onnx.symbolic_opset17
107.. py:module:: torch.onnx.symbolic_opset18
108.. py:module:: torch.onnx.symbolic_opset19
109.. py:module:: torch.onnx.symbolic_opset20
110.. py:module:: torch.onnx.symbolic_opset7
111.. py:module:: torch.onnx.symbolic_opset8
112.. py:module:: torch.onnx.symbolic_opset9
113.. py:module:: torch.onnx.utils
114.. py:module:: torch.onnx.verification
115.. py:module:: torch.onnx.symbolic_opset12