...

Problem

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

...

{
    "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]
              }
}

...

Possible use cases

Inputs and outputs are dictionary where key => Name of the node and value => shape. Though single input, single output models are more common, a model can be of following variations: 

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

User Experience

Before

Step 1 - Train and export the model from Gluon

# Trained network
net = mx.gluon.model_zoo.vision.resnet18_v1(pretrained=True, ctx=mx.cpu())

# Data transformations applicable during inference
inference_input_transforms = gluon.nn.HybridSequential()
inference_input_transforms.add(transforms.Resize((224, 224)))
inference_input_transforms.add(transforms.ToTensor())
inference_input_transforms.add(transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)))

# Export the model. Cannot export data transformation and input/output signature

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

...

Step 2 - Import the model for inference with Module OR Gluon OR JAVA

Import the model in Gluon

### Load netthe =>model. HybridModel Sequentialdoes Network
## input_transforms => Hybrid Sequential Transformation Pipeline
gluon.contrib.utils.export(net, path="./my_model", not contain transformation and input/output signature 
net = SymbolBlock.imports(symbol_file="my_model-symbol.json",
                          input_names=["data"],
                           epoch=0, param_file="my_model-0000.params",
                           signature={constants.INPUT_DESC:[("data"ctx=mx.cpu())

Import the model in Module API

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

Import the model in Java Predictor API

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

Step 1 - Train and export the model from Gluon

# Trained network
net = mx.gluon.model_zoo.vision.resnet18_v1(pretrained=True, ctx=mx.cpu())

# Data transformations applicable during inference
inference_input_transforms  constants.OUTPUT_DESC:[("softmax_label", (1,10))]},= gluon.nn.HybridSequential()
inference_input_transforms.add(transforms.Resize((224, 224)))
inference_input_transforms.add(transforms.ToTensor())
inference_input_transforms.add(transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)))

# Export the model
gluon.contrib.utils.export(net, path="./my_model", 
                           input_transforms={"data":input_transforms},epoch=0, 
                           output_transforms=None)

Import API - Gluon

Before

SymbolBlock.imports(symbol_file="my_model-symbol.json",
signature={constants.INPUT_DESC:[("data", (1,3,224,224))],
                         input_names=["data"]             constants.OUTPUT_DESC:[("softmax_label", (1,10))]},
                    param_file="my_model-0000.params"       input_transforms={"data":inference_input_transforms},
                    ctx='cpu')

...

       output_transforms=None)

Step 2 - Import the model for inference with Module OR Gluon OR JAVA

Import the Model in Gluon

# Load the model along with the transformations
net = gluon.contrib.utils.import(symbol_file="my_model-symbol.json",
                           		 param_file="my_model-0000.params",
                           		 load_transforms = True,
                           		 ctx = 'cpu')


# Prediction
pred = net(data)

Import the Model in Module API

...

(Supported to create a module for inference only)Before

Code Block
language py
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.

Code Block
language py
# Load the model along with the transformations. Can be used only for inference (forward()) 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) # Prediction mod.forward(...)

Import

...

the Model in Java Predictor API

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-my_model";


# Load the model along with the transformations.
Predictor predictor = new Predictor(modelPathPrefix, inputDescListcontext, contextload_transforms=True);

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

After

...

...

Performance Consideration

During inference, initial benchmarks shows a noticeable performance gain with End to end models i.e., a model with data transformations and neural network all fused as a single model graph.

• ResNet-18 model pre-trained with ImageNet. https://s3.us-east-2.amazonaws.com/mxnet-public/end_to_end_models
• Pre-processing - Resize(224, 224), ToTensor, Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
• We take average of 500 runs
• Single Request Inference - Input Data - Synthetic (random.uniform(0, 255, shape=(1, 300, 300, 3))
• Batch Inference - Input Data - Synthetic (random.uniform(0, 255, shape=(25, 300, 300, 3))
• Below time gives - Average Prediction Time Per Sample

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. If a user tries to load new MXNet model with older MXNet version, they get errors such as - "unknown field 'inputs', 'outputs' in the model" because, MXNet's model JSON parser schema in old MXNet do not understand new fields introduced as part of this work.

Open Questions and Assumptions Made

1. Here we assume internal operators as no op connector node. We assume these connector nodes are not exposed to the users and hence, we can safely use as connector nodes for fusing transformations and neural network graph. Is this assumption in the right direction?
2. Is this ok to go to Contrib first and then update main APIs? Below are the reasoning:
   1. API changes / additions in this work is proposed to go in to Contrib because "export" and "import" APIs are most widely used by every MXNet users. We propose to mature the API, address any usability or design concerns while in contrib before graduating to main APIs.
   2. What is the plan to monitor and move this change from contrib to main? Once we build these extensions/new export, import APIs, we will be creating new tutorials and update the existing export/import tutorials. We will track issues and discussion for any user concerns. We will graduate these APIs after maturing in Contrib for at least 1 version of MXNet release.

Alternate Solution

1. Keep transformation graph and network graph separately independent of each other and fuse them at run time. 
   1. 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
   2. In this alternate solution, we 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.
   3. Two concerns for this approach:
      1. In MXNet, a symbol and a param are foundational building block of representing 1 graph and its parameters. It may not be ideal to keep multiple graphs in the same symbol file and params file.
      2. If we decide to keep transformations and network graphs as separate files, this can be still achieved today with current MXNet APIs without any additional changes.

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