Problems

Taking a MXNet model from training and deploying it in production poses several challenges to the users. Most important problems raised by users are:

1. Input/Output data transformations are not part of the MXNet model.
2. Input/Output signature are not part of the model: Saved model missing the information about the input/output descriptions, like name/shape, making the saved model unusable out of the box.
3. File naming - Epoch number is part of model file name. This information may not be necessary for production deployment.

Proposed Solution

1. In Gluon, update the Export API to accept input/output signature, pre-processing/post-processing transformations (Hybrid Sequential Block) from the users.
2. If user provides the transformations, a fused graph (Transformations + Batchify_Connector_Node + Network + No_Op_Connector_Node + Transformations) is prepared. Batchify_Connector_Node and No_Op_Connector_Node are MXNet internal identifier operators to help in identifying transformations and network graphs separately. This identifier nodes will be helpful, if when loading the model, user prefers to not load the transformations.
3. Input/Output signature are added to the symbol file.

Example symbol file after proposed update update. "inputs" and "outputs" are the new parameters proposed in this work.

{
    "nodes": [..........]
    "arg_nodes": [0, 1, 2, 4, 5],
    "node_row_ptr": [0, 1, 2, 3, 4, 5, 6, 7, 8],
    "heads": [[7, 0, 0]],
    "attrs": {"mxnet_version": ["int", 10301]
              "inputs": {"data":[1,3,224,224]},
              "outputs" : {"softmax_label":[1,10]
              }
}

Use cases

1. Single input, single output
2. Single input, multiple outputs
3. Multiple input with same or different transformation, single output
4. Multiple input with same or different transformation, multiple output

Export API - Gluon

1. Since this is an update to the most commonly used API - Export in Gluon, we will first introduce the change with a new Export API in gluon.contrib.

Before

net.export(path="./my_model", epoch=0)

After

## net => Hybrid Sequential Network
## input_transforms => Hybrid Sequential Transformation Pipeline
gluon.contrib.utils.export(net, path="./my_model", 
                           epoch=0, 
                           signature={constants.INPUT_DESC:[("data", (1,3,224,224))],
                                      constants.OUTPUT_DESC:[("softmax_label", (1,10))]},
                           input_transforms={"data":input_transforms},
                           output_transforms=None)

Import API - Gluon

Before

SymbolBlock.imports(symbol_file="my_model-symbol.json",
                    input_names=["data"],
                    param_file="my_model-0000.params",
                    ctx='cpu')

After

gluon.contrib.utils.import(symbol_file="my_model-symbol.json",
                           param_file="my_model-0000.params",
                           load_transforms = True,
                           ctx = 'cpu')

Import API - Module

Before

sym, arg_params, aux_params = mx.model.load_checkpoint('my_model', 0)
mod = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
mod.bind(for_training=False, data_shapes=[('data', (1,3,224,224))], 
         label_shapes=mod._label_shapes)
mod.set_params(arg_params, aux_params, allow_missing=True)

mod.forward(...)

After
Supported to create a module for inference only.

mod = mx.contrib.Module.import(
                symbol_file = "my_model-symbol.json",
                param_file = "my_model-0000.params",
                load_transforms = True,
                ctx = 'cpu',
                batch_size = 1)
mod.forward(...) 

Import API - Java predictor

Before

Shape inputShape = new Shape(new int[] {1,3,224,224});
DataDesc inputDescriptor = new DataDesc("data", inputShape, DType.Float32(), "NCHW"); 
List<DataDesc> inputDescList = new ArrayList<DataDesc>();
inputDescList.add(inputDescriptor);
List<Context> context = new ArrayList<>();
context.add(Context.cpu()); 
String modelPathPrefix = "path-to-model";
Predictor predictor = new Predictor(modelPathPrefix, inputDescList, context);

List<NDArray> result = predictor.predictWithNDArray(inputNDArray);

After

List<Context> context = new ArrayList<>();
context.add(Context.cpu()); 
String modelPathPrefix = "my_model";
Predictor predictor = new Predictor(modelPathPrefix, context, load_transforms=True);

List<NDArray> result = predictor.predictWithNDArray(inputNDArray);

Backward compatibility

1. All APIs changes are backward compatible.
2. Old MXNet model should still load without breakage with new MXNet version.
3. New MXNet model will not work on old MXNet versions.

Alternate Solution

1. Keep transformation graph and network graph separately independent of each other and fuse them at run time. In the proposed approach, we fuse the transformation and neural network and export as single graph. We introduce a no_op_identifier operator to identify the link between transformation and neural network. Another solution would be to keep the transformation and network graph separately (in same symbol file or multiple symbol file). These independent graphs can then be fused during run time.

References

1. Gluon-CV export helper - https://gluon-cv.mxnet.io/api/utils.html?highlight=export#gluoncv.utils.export_block
2. https://mxnet.incubator.apache.org/versions/master/tutorials/python/predict_image.html?highlight=load_checkpoint
3. https://medium.com/apache-mxnet/introducing-java-apis-for-deep-learning-inference-with-apache-mxnet-8406a698fa5a