ONNX_MXNet API Design

This document describes the onnx module in mxnet.contrib package that provides ONNX format support within MXNet. It outlines the currently implemented APIs and the future roadmap and design of proposed APIs.

Use cases for the module:

1. Import ONNX models into MXNet symbolic interface.
2. Import ONNX model into MXNet Gluon interface.
3. Export MXNet symbolic model to an ONNX model file. (TODO)
4. Export MXNet Gluon model into an ONNX model file.(TODO)

To meet the above use cases, we design the following APIs –

• Import ONNX into MXNet Symbol graph

Given an ONNX model file import into MXNet’s symbolic graph along with all the parameter tensors. This API is implemented and will be shipped as part of MXNet v1.2 release.

The API definition is as follows –

mxnet.contrib.onnx.import_model(model_file) 

Imports the ONNX model file, passed as a parameter, into MXNet symbol and parameters. Operator support and coverage - https://cwiki.apache.org/confluence/display/MXNET/ONNX

Input Parameters –

model_file – A string object representing the path to the ONNX model file.

Return type –

sym – An mxnet symbol object representing the symbolic graph of the given model.

arg_params -  A dictionary object mapping the parameter name to an mxnet ndarray object representing its tensor value. These are the parameter values that are learned while training the model.

aux_params – A dictionary object mapping parameter names to an mxnet ndarray object representing the tensor values of the parameters. It stores the values that are not learned during the training process.

• [Proposed API] Import ONNX model files into Gluon Symbolic block

Given an ONNX model file import the model into Gluon’s SymbolicBlock object.

mxnet.contrib.onnx.import_to_gluon(model_file)

Imports the ONNX model files, passed as a parameter, into Gluon’s SymbolicBlock object.

Input Parameters –

model_file – A string object representing the path to the ONNX model file.

Output Parameters –

sym_block – A SymbloicBlock object representing the given model file.

• [Proposed API] Model Metadata for a given ONNX model file

Given an ONNX model file, the user can use this API to fetch the related metadata of the model. This is a request from customers and users of the ONNX module, where they had a use case for knowing the shape information of the input and output tensors of a given ONNX model. This API is being implemented to meet that request.

mxnet.contrib.onnx.get_model_metadata(model_file)

Returns the metadata information of the given model. Currently it will return the shape information of the input and output tensors.

Input Parameters –

model_file – A string object representing the path to the ONNX model file.

Return Type –

model_metadata – A dictionary object mapping various metadata to its corresponding value. The proposed implementation will return a dictionary with following template.

{

“input_tensor_shape”: <list of tuples representing the shape of the input paramters>,

“output_tensor_shape”: <list of tuples representing the shape of the output of the model>

}