torch.onnx

Overview

Open Neural Network eXchange (ONNX) is an open standard format for representing machine learning models. The torch.onnx module captures the computation graph from a native PyTorch torch.nn.Module model and converts it into an ONNX graph.

The exported model can be consumed by any of the many runtimes that support ONNX, including Microsoft’s ONNX Runtime.

There are two flavors of ONNX exporter API that you can use, as listed below. Both can be called through function torch.onnx.export(). Next example shows how to export a simple model.

import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 128, 5)

    def forward(self, x):
        return torch.relu(self.conv1(x))

input_tensor = torch.rand((1, 1, 128, 128), dtype=torch.float32)

model = MyModel()

torch.onnx.export(
    model,                  # model to export
    (input_tensor,),        # inputs of the model,
    "my_model.onnx",        # filename of the ONNX model
    input_names=["input"],  # Rename inputs for the ONNX model
    dynamo=True             # True or False to select the exporter to use
)

Next sections introduces the two versions of the exporter.

TorchDynamo-based ONNX Exporter

The TorchDynamo-based ONNX exporter is the newest (and Beta) exporter for PyTorch 2.1 and newer

TorchDynamo engine is leveraged to hook into Python’s frame evaluation API and dynamically rewrite its bytecode into an FX Graph. The resulting FX Graph is then polished before it is finally translated into an ONNX graph.

The main advantage of this approach is that the FX graph is captured using bytecode analysis that preserves the dynamic nature of the model instead of using traditional static tracing techniques.

Learn more about the TorchDynamo-based ONNX Exporter

TorchScript-based ONNX Exporter

The TorchScript-based ONNX exporter is available since PyTorch 1.2.0

TorchScript is leveraged to trace (through torch.jit.trace()) the model and capture a static computation graph.

As a consequence, the resulting graph has a couple limitations:

• It does not record any control-flow, like if-statements or loops;

• Does not handle nuances between training and eval mode;

• Does not truly handle dynamic inputs

As an attempt to support the static tracing limitations, the exporter also supports TorchScript scripting (through torch.jit.script()), which adds support for data-dependent control-flow, for example. However, TorchScript itself is a subset of the Python language, so not all features in Python are supported, such as in-place operations.

Learn more about the TorchScript-based ONNX Exporter

Contributing / Developing

The ONNX exporter is a community project and we welcome contributions. We follow the PyTorch guidelines for contributions, but you might also be interested in reading our development wiki.