...

To implement the import/export functionality in MXNet, I propose to expose a MXNet python module as a wrapperto do the model conversion:

• “serde”(serialization/deserialization - name taken from Apache Hive project) which can be called as (mx.serde.import/ mx.serde.export)
• OR It can also go under contrib. MXNet contrib is used to add experimental features which can be later moved outside. This can qualify as an experimental feature as there is a gap in operator implementation. (See Appendix)

 mx.contrib.serde.import/ mx.contrib.serde.export

This wrapper serialization/deserialization module should have following methods supporting different formats:

Import model into mxnet.

sym, params = mx.serde.import(input_file, ininput_format=‘onnx’, outoutput_format=’gluon’)

• input_file : input model file (e.g., protobuf model file for onnx)
• ininput_format : (optional) onnx, coreml
• outoutput_format : (optional) “gluon/symbolic”. by default, gluon will be used.

...

Note: Currently gluon does not provide an easy way to import a pre-trained model. (there is a workaround using which this can be done). 

Export mxnet model to specified input format.

1) mx.serde.export(sym, params, input_format='symbolic', output_format=‘onnx’, filename_prefix=”model_name”)

2) mx.serde.export(module, input_format='symbolic', output_format=‘onnx’, filename_prefix=”model_name”)

3) mx.serde.export(gluon_model, input_format='gluon', output_format=‘onnx’, filename_prefix=”model_name”)

• sym : model definition
• module : mxnet module object
• gluon_model : model definition (HybridBlock)
• params : weights
• input_format : symbolic/gluon 
• output_format : onnx onnx, coreml
• filename_prefix: a filename prefix to be used to model save files. E.g., for onnx, a binary protobuf will be written to output file with “.onnx” extension.

• files. E.g., for onnx, a binary protobuf will be written to output file with “.onnx” extension.

In addition to this core format conversion module(`serde`), I think implementing a wrapper according to the different namespace will be more user friendly.

For example,

•  Gluon

mx.gluon.import(input_file, input_format=‘onnx’)

mx.gluon.export(output_format=’onnx’, filename_prefix=”model_name”)

If any pre-processing/post-processing logic is required for gluon specific models, it can go under this gluon wrapper. These functions will internally call `serde` module APIs.

e.g. `mx.gluon.import(input_file, input_format=‘onnx’)` will internally call `mx.serde.import(input_file, input_format=‘onnx’, output_format='gluon')`

•  For Symbolic interface

sym, params = mx.mod.Module.import(input_file, input_format=‘onnx’)

mx.mod.Module.export(output_format=‘onnx’, filename_prefix=”model_name”)

This function will directly save the file called “model_name.onnx” on disk. 

Implementation approaches

There are two approaches which can be taken to import/export onnx model.

...

•   Through MXNet's symbolic operators

Implement at the MXNet layer by parsing the ONNX model(in protobuf format) and turn into MXNet Symbolic operators and build MXNet model directly. Similarly, MXNet model can be converted to ONNX format at this layer.

...

The whole implementation will go under MXNet repo.

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

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

...

•   Through

...

• NNVM/top operators

  
  

The DMLC community has released the nnvm/tvm complier and an intermediate representation of the models.

...

import nnvm.frontend

def import(..., input_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(..., output_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, I propose to take the 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/TVM package.

...

data_iter = data_process(mnist)
sym = construct_model_def(...)
# create module for training
mod = mx.mod.Module(symbol=sym, ...)
# train model
mod.fit(data_iter, optimizer='...', optimizer_params='...', ...)
# get parameters of trained model
params = mod.get_params()

# save into different format using sym/params

# This will internally call serde package API `serde.export`
mx.mod.serdeModule.export(sym, params, output_format=‘onnx’, filename=”model.onnx”)

# OR save into different format using module directly

mx.serdemod.Module.export(mod, output_format=‘onnx’, filename=”model.onnx”)

...

The general structure of onnx import would look something like this:

# Import model

# This will internally call serde package API `serde.import`

sym, params = mxnet.mod.serdeModule.import(input_file, ininput_format=‘onnx’)

# create module for inference
mod = mx.mod.Module(symbol=sym, data_names=..., context=mx.cpu(), label_names=None)
mod.bind(for_training=False, data_shapes=..., label_shapes=None)

# set parameters
mod.set_params(arg_params=params, aux_params=None, allow_missing=True)

# forward on the provided data batch
mod.forward(Batch([data_batch]))

...