Machine Learning / Software Development

Tutorial: Import an ONNX Model into TensorFlow for Inference

24 Jul 2020 9:22am, by
This post is the fourth in a series of introductory tutorials on the Open Neural Network Exchange (ONNX), an initiative from AWS, Microsoft, and Facebook to define a standard for interoperability across machine learning platforms. See: Part 1, Part 2, and Part 3.

In the last tutorial, we trained a CNN model in PyTorch and converted that into an ONNX model. In the current tutorial, we will import the model into TensorFlow and use it for inference.

Before proceeding, make sure that you completed the previous tutorial as this is an extension of the same.

Converting ONNX Model to TensorFlow Model

The output folder has an ONNX model which we will convert into TensorFlow format.

ONNX has a Python module that loads the model and saves it into the TensorFlow graph.

We are now ready for conversion. Create a Python program with the below code and run it:

The output folder contains three models: PyTorch, ONNX, and TensorFlow.

We are now ready to use the model in TensorFlow. Note that it works only with TensorFlow 1.x. For this tutorial, we are using the 1.15, which is the last version.

We start by importing the right modules and then disable the warnings generated by TensorFlow.

The names for input and output tensor can be taken from Netron tool by opening the model.pb file.

The input node (input.1) and output node (add_4) name and shape are visible in the Netron.

The next few lines of code preprocess the image through OpenCV. We then open the TensorFlow model and create a session based on the graph.

Finally, by applying the argmax function, we classify the output into one of the ten classes defined by MNIST.

In this tutorial, we imported an ONNX model into TensorFlow and used it for inference. In the next part, we will build a computer vision application that runs at the edge powered by Intel’s Movidius Neural Compute Stick. The model uses an ONNX Runtime execution provider optimized for the OpenVINO Toolkit. Stay tuned.

Janakiram MSV’s Webinar series, “Machine Intelligence and Modern Infrastructure (MI2)” offers informative and insightful sessions covering cutting-edge technologies. Sign up for the upcoming MI2 webinar at

Feature image: Second-order Cauchy stress tensor, Wikipedia.