Best Practices for Neural Network Exports to ONNX

Photo by Artem Sapegin on Unsplash.

ONNX is an open format built to represent machine learning models. ONNX defines a common set of operators — the building blocks of machine learning and deep learning models — and a common file format to enable AI developers to use models with a variety of frameworks, tools, runtimes, and compilers.

— onnx.ai

Exporting your model to ONNX helps you to decouple the (trained) model from the rest of your project. Moreover, exporting also avoids environment dependencies like on the python interpreter, framework version, used libraries, etc. And the exported ONNX-model can store both, the architecture and parameters of your model. This means that it is enough to send your co-workers a single file to exchange your model.

Our experience shows that is easier to export PyTorch models. If possible, choose a PyTorch source and convert it using the built-in torch.onnx module. Alternatively, you can use the newer standalone onnx python package (just replace torch.onnx with onnx in the following code examples).

from PyTorch

PyTorch models can only be exported programmatically:

Please note that PyTorch computes the computation graph at runtime, therefore the conversion engine requires a batch of correct shape (the data can be random in most cases) that it will pass through the network to understand the architecture. torch.onnx uses torch.jit.trace to find the route your data takes through the network.

Best practices:

• Watch out for TracerWarnings. These may indicate that the tracer could not follow your batch.
• If you make routing decisions at runtime, make sure to use a batch/configuration that takes all routes you are trying to export.
• If you need to make routing decisions you should make them on Tensors. Pythons default types are not found by the tracer.
• If you are using torchhub models, check if they provide a exportable parameter (or similar) to replace incompatible operations.

from TensorFlow (1/2/lite/js/tf.Keras)

We recommend Microsofts tf2onnx package for the conversion of TensorFlow models. You have to store the model to one of TensorFlows supported file formats prior to the ONNX export. Supported formats include saved model, checkpoint, graphdef or tflite.

Export savedmodel files to ONNX:

python -m tf2onnx.convert --saved-model tensorflow-model-path --output model.onnx

This is how you do it for tflite (alternatively use tflite2onnx):

python -m tf2onnx.convert --opset 13 --tflite tflite--file --output model.onnx

For the other formats, you need to provide the names of the input and output tensors. tf2onnx will use them to trace the network. Providing wrong or not a complete list of labels may result in a corrupted export. If you do not know the input and output node names of the model, you can use the summarize_graph TensorFlow utility. Here is how you can install and use it:

Alternatively, check the source project from prominent placement or ask the original author. For the following examples, we assume that there are two tensorflow inputs named input0:0,input1:0 and one output named output0:0.

For checkpoint format:

python -m tf2onnx.convert --checkpoint tensorflow-model-meta-file-path --output model.onnx --inputs input0:0,input1:0 --outputs output0:0

For graphdef format:

python -m tf2onnx.convert --graphdef tensorflow-model-graphdef-file --output model.onnx --inputs input0:0,input1:0 --outputs output0:0

Note: Models export for reuse (such as TensorFlow Hub models) are not analyzable with summarize_graphs and may not be exportable at all.

The ONNX protocol, or, the used converter may not support all operations of the source model.

Possible solutions:

• Check if there is a newer operation set version (opset) available that supports your operation in question.
• Check if the non-supported operation can be replaced by a supported one e.g. it is common to replace Hardswish or SiLU activations with ReLU (as of opset 11).
• If your converter simply does not map the operation but the operation is supported by the ONNX protocol implement the mapping or try a different converter.
• If the operation is not supported by the ONNX protocol try to either rewrite your operation with supported operations or implement a mapping using supported operations. Also, consider submitting a PR to ONNX.

We recommend loading the model using PyTorch and using the built-in validation engine.

This code block will only validate the schema. This does not guarantee that your architecture is complete, nor that all parameters were (correctly) exported. Consequently, we advise you to run inference on a few samples and compare them to inference with your original framework. Note that there may be a slight difference due to the export process and a potentially different execution framework. To use inference make sure to install the onnxruntime python package with pip or your python package manager.

Additionally, you can run a sanity-check by visualizing the exported model with an external tool like Netron. Netron also allows you to explore stored parameters. Keep in mind that the converters may reduce or expand operations as needed, hence, are potentially touching up the original architecture. Note: Netron may have issues loading large models.

Whether you choose ONNX or not — always consider publishing your trained models, especially if you are preparing a scientific publication. This helps others to easily replicate your findings and also provides a great start for their projects.