5min: Gluon Model to ONNX

Exporting Gluon model to ONNX format via MXNet Symbol.

Gluon API provides an easy way to prototype, build & train deep learning models.

Apache MXNet, the deep learning framework, has always been known to have that rare combination of being easy for research & best suited for enterprise/production use-cases.

Gluon provides a very easy interface for leveraging MXNet. The straight dope book provides excellent material for getting up-to-speed with Gluon API & deep learning in general.

ONNX.ai, Open Neural Network Exchange, is an open ecosystem for interoperable AI models. onnx/onnx [Github repo] acts as an open standard for machine learning interoperability i.e. allows models from different frameworks to be converted into one another.

Gluon defines neural networks in terms of blocks: Block, HybridBlock, SymbolBlock, etc. An easy way to convert your gluon model to ONNX is via MXNet Symbol object.

Step 0: Setup Environment

Create a virtual environment using Python Package virtualenv. Virtual Environment ensures isolated environment which is perfect for our use-case. Read more on why to create virtual environments.

# create & use python package: virtualenv
virtualenv -p python3 test_onnx_13_mx
source test_onnx_13_mx/bin/activate

# install onnx & mxnet
pip install onnx==1.3.0
pip install mxnet
pip install gluoncv

Step 1: Download Gluon Model

from gluoncv import model_zoo
model_name = 'resnet50_v1'
resnet50 = model_zoo.get_model(model_name, pretrained=True)                       print(model_name+' downloaded')

Step 2: Convert Gluon model to MXNet Symbol object

To convert the loaded gluon model into MXNet’s Symbol object, we have to hybridize the model. Hybridization allows users to leverage advantages of imperative code style and symbolic mode performance. It’s like “best of both worlds”.

Hybridization is the best of both worlds, Symbolic & Imperative.

resnet50.hybridize()                       
print(model_name+' hybridized')

After hybridization, the model still needs to go through atleast one pass of data to be able to store the hybrid model into a symbol. As a result, we create input data of the input shape which is compatible to the model. Resnet50 takes input of shape (3, 224, 224) based on CHW format.

C : channel [RGB]
H : height
W : width

For the purpose of this experiment, we can pass the batch_size as 1.

input_shape=(1,3,224,224)
data_array = np.random.uniform(0,255,size=input_shape).astype("float32")
mx_data = mx.nd.array(data_array)
resnet50(mx_data)

Step 3: Export model

The endgoal of model training is to deploy it into production infrastructure or store it for reproducibility. And the way to ensure the model is deploy-able is by exporting it.

With Gluon, network parameters are associated with a Block, but the execution flow is controlled in python through the code in .forward() function. Hence only hybridized networks can be exported with a -symbol.json and .params file using .export(), non-hybridized models can only have their parameters exported using .save_parameters(). Check this great tutorial to learn more: Saving and Loading Gluon Models.

resnet50.export(model_name)
print(model_name+' exported')

Step 4: Convert to ONNX

Once the model is saved in the form of symbol.json and .params files, we can safely convert it to ONNX using the following commands

#convert using onnx
from mxnet.contrib import onnx as onnx_mxnet
onnx_file='./output.onnx'
params = './'+model_name+'-0000.params'
sym='./'+model_name+'-symbol.json'
onnx_mxnet.export_model(sym, params, [input_shape], np.float32, onnx_file)
print('onnx export done')

That’s it! You’ve a model imported from Gluon, converted to MXNet Symbol and finally into ONNX standard model format.

To access the entire script: