Convert ONNX Format to TensorFlow 2 Model

ONNX is a format for representing machine learning models. ONNX is like an intermediary that makes it easier to exchange models between different machine learning frameworks. This tutorial demonstrates how to convert a model represented in ONNX format to TensorFlow 2 model.

Prepare environment

Before starting, install the following packages:

pip install tensorflow

pip install tensorflow-probability

pip install onnx-tf

The tensorflow-probability package is required by onnx-tf which is used for model conversion.

Model training

Read the separate post how to convert TensorFlow 2 model to model represented in ONNX format. Once you have the model.onnx file, continue this tutorial.

Model conversion

Load the .onnx file. Convert model represented in ONNX format to model in SavedModel format, which can be loaded in TensorFlow 2.

convert.py

import onnx
from onnx_tf.backend import prepare

model = onnx.load('model.onnx')

prepare(model).export_graph('model')

Inference

Load SavedModel using a low-level tf.saved_model.load function. Then, predict a value of y for a previously unknown value of x.

test.py

import tensorflow as tf

model = tf.saved_model.load('model')
infer = model.signatures['serving_default']

x = tf.constant([[15.0]])

outputs = list(infer.structured_outputs)
y = infer(x)[outputs[0]].numpy()
print(y)

In our case, y is equal to 31.00175 when x is 15.0. Verify result:

y = 2 * x + 1 = 2 * 15 + 1 = 31