ONNX Model: Export Using Pytorch, Problems, and Solutions

As we know that there a lot of AI or ML frameworks available in the market (paid or free version) to develop your AI/ML solutions. But the major problem is they don’t provide support within themselves. Hence, if you develop a model in Caffe(an AI/ML development framework), it won’t be possible for you to use the same model in Pytorch and vise versa.

Thanks to the massive research community demand and the need for a system to switch between different frameworks. In September 2017 Facebook and Microsoft introduced a system for switching between machine learning frameworks such as PyTorch and Caffe2. Later, IBM, Huawei, Intel, AMD, ARM, and Qualcomm announced support for the initiative.

What is Open Neural Network Exchange (ONNX)?

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 is available on GitHub, where it defines itself as an open ecosystem that empowers AI developers to choose the right tools as their project evolves. Currently, it focuses on the capabilities needed for inferencing (scoring).

The connection between ONNX and TensorRT

As it’s clear that ONNX facilitates a common file format for the AI developers (like me ??) to code in any framework(Pytorch in my case) and then export the trained model in ONNX format and use it for TensorRT inference. TensorRT provides ONNX support by ONNX model parser and “onnx_graphsurgeon” library.

Problem with ONNX, Pytorch, and TensorRT parsing.

1. Develop and train a model using Pytorch. (Only PyTorch 1.4 or above is supported in TensorRT)
2. Export PyTorch trained model in ONNX format. (Older ONNX format is not supported in TensorRT. Also, many mathematical operations are still not supported in the ONNX [ongoing development])
3.  Dynamic Batching and dynamic convolution are still a problem.
4. TensorRT ONNX parser fails in parsing complex graphs of the exported model. (This part is very interesting and a separate article is dedicated to this in this article series)

Here are some of the solutions:

1. Use one version lower to that of the latest stable version of Pytorch (for example current stable release is 1.7. Then use 1.6). Don’t follow the instructions of using very older versions, as it may cause conflict in operations while exporting the model in ONNX.
2. Always, use the latest ONNX while exporting the model. Also, always try to use the latest opset, for example, the current latest is "opset11". The latest Opset allows better export of the model graph.
3. Dynamic batching can be achieved in Pytorch and it’s very easy. Here is code example and tutorial for the same.

Dynamic convolution or dynamic image size is still under development and expected to be part of the next release.

Articles in this series"

1- Model optimization for Fast Inference and Quantization

2- TensorRT Installation: Problems and "Way Arounds"

 Failing of the ONNX model parsing in tensroRT is associated with the complex model graph and how to deal with it is the part of the coming article. Also, we will see that how we can manually set the dynamic batching to a ONNX model using ONNX library

Till then enjoy oranges (mandarini)…!!

thanks :)