[v1.x] ONNX export support from arg/aux params

python/mxnet/onnx/mx2onnx/_export_model.py

@@ -55,14 +55,16 @@ def export_model(sym, params, in_shapes=None, in_types=np.float32,
     """Exports the MXNet model file, passed as a parameter, into ONNX model.
     Accepts both symbol,parameter objects as well as json and params filepaths as input.
     Operator support and coverage -
-    https://cwiki.apache.org/confluence/display/MXNET/ONNX+Operator+Coverage
+    https://github.com/apache/incubator-mxnet/tree/v1.x/python/mxnet/onnx#operator-support-matrix
 
     Parameters
     ----------
     sym : str or symbol object
         Path to the json file or Symbol object
-    params : str or symbol object
-        Path to the params file or params dictionary. (Including both arg_params and aux_params)
+    params : str or dict or list of dict
+        str - Path to the params file
+        dict - params dictionary (Including both arg_params and aux_params)
+        list - list of length 2 that contains arg_params and aux_params
     in_shapes : List of tuple
         Input shape of the model e.g [(1,3,224,224)]
     in_types : data type or list of data types
@@ -89,7 +91,7 @@ def export_model(sym, params, in_shapes=None, in_types=np.float32,
 
     Notes
     -----
-    This method is available when you ``import mxnet.contrib.onnx``
+    This method is available when you ``import mxnet.onnx``
 
     """
 
@@ -126,6 +128,15 @@ def export_model(sym, params, in_shapes=None, in_types=np.float32,
                                                        in_types_t,
                                                        verbose=verbose, opset_version=opset_version,
                                                        dynamic=dynamic, dynamic_input_shapes=dynamic_input_shapes)
+    elif isinstance(sym, symbol.Symbol) and isinstance(params, list) and len(params) == 2:
+        # when params contains arg_params and aux_params
+        p = {}
+        p.update(params[0])
+        p.update(params[1])
+        onnx_graph = converter.create_onnx_graph_proto(sym, p, in_shapes,
+                                                       in_types_t,
+                                                       verbose=verbose, opset_version=opset_version,
+                                                       dynamic=dynamic, dynamic_input_shapes=dynamic_input_shapes)
     else:
         raise ValueError("Input sym and params should either be files or objects")
 

tests/python-pytest/onnx/test_onnxruntime_cv.py

@@ -61,6 +61,30 @@ def export_onnx_dynamic(self, dynamic_input_shapes):
                              dynamic_input_shapes=dynamic_input_shapes)
         return onnx_file
 
+    def export_onnx_argaux(self):
+        onnx_file = self.modelpath + ".onnx"
+        sym_file = self.modelpath + "-symbol.json"
+        params_file = self.modelpath + "-0000.params"
+        if not (os.path.isfile(sym_file) and os.path.isfile(params_file)):
+            raise ValueError("Symbol and params files provided are invalid")
+
+        try:
+            # reads symbol.json file from given path and
+            # retrieves model prefix and number of epochs
+            model_name = sym_file.rsplit('.', 1)[0].rsplit('-', 1)[0]
+            params_file_list = params_file.rsplit('.', 1)[0].rsplit('-', 1)
+            # Setting num_epochs to 0 if not present in filename
+            num_epochs = 0 if len(params_file_list) == 1 else int(params_file_list[1])
+        except IndexError:
+            logging.info("Model and params name should be in format: "
+                         "prefix-symbol.json, prefix-epoch.params")
+            raise
+
+        sym, arg_params, aux_params = mx.model.load_checkpoint(model_name, num_epochs)
+        params = [arg_params, aux_params]
+        mx.onnx.export_model(sym, params, [self.input_shape], self.input_dtype, onnx_file)
+        return onnx_file
+
     def predict(self, data):
         return self.model(data)
 
@@ -176,7 +200,11 @@ def normalize_image(imgfile):
     try:
         tmp_path = str(tmp_path)
         M = GluonModel(model, (1,3,inlen,inlen), 'float32', tmp_path)
-        onnx_file = M.export_onnx()
+        if model == 'resnet50_v2':
+            # testing export for arg/aux
+            onnx_file = M.export_onnx_argaux()
+        else:
+            onnx_file = M.export_onnx()
 
         # create onnxruntime session using the generated onnx file
         ses_opt = onnxruntime.SessionOptions()