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.
...