This document proposes a new module in MXNet to support import/export functionalities for different formats like ONNX, CoreML.

Table of Contents

Objectives

ONNX is an intermediate representation for describing a neural network computation graph and weights. With AWS, Microsoft and Facebook defining and promoting ONNX, major Deep Learning frameworks such as MXNet, Pytorch, Caffe2, and CNTK are building native support for model import and export.

...

4)    Export MXNet Gluon model into ONNX. 

Future Use Cases (out of scope of this design):

...

• In the future, we may have to reimplement at the nnvm/tvm layer, in case MXNet moves to the nnvm/tvm backend. If this happens, we will need to implement conversion for the existing symbolic operators for backward compatibility which we can leverage for onnx-mxnet conversion as well.
• MXNet's API contains some legacy issues which are supposedly fixed in nnvm/top operators. #issue lists down some of the issues and plan to fix them.

Internal API design

Whole The whole implementation will go under MXNet repo.

def import(..., format='onnx'):
    # returns mxnet graph with parameters
    ...
    return sym, params

def export(..., format='onnx'):
  # returns ONNX protobuf object
  ...
  return onnx_proto

1)   Through nnvm/top operators

...

Wrapper API(serde) will go under mxnet repo which will internally call nnvm package methods.

import nnvm.frontend

def import(..., format='onnx'):
  # convert from onnx to nnvm graph
  nnvm_graph, params = nnvm.frontend.from_onnx(...) # Exists
  # convert fron nnvm graph to mxnet graph
  mxnet_graph, params = nnvm.frontend.to_mxnet(...) # Need to implement
  return mxnet_graph, params

def export(..., format='onnx'):
  # convert from mxnet to nnvm graph
  nnvm_graph, params = nnvm.frontend.from_mxnet(...) # Exists
  # convert fron nnvm graph to onnx proto format
  onnx_proto = nnvm.frontend.to_onnx(...) # Need to implement
  return onnx_proto

Suggested approach:

As a middle ground for both of the above implementation choices, we can take first approach and implement MXNet->ONNX conversion for export functionality and if someone wants to take advantage of NNVM/TVM optimized engine for their usage, they can do it by leveraging import functionality provided in NNVM package.

...