[incubator-mxnet] branch v1.x updated: [v1.x] More onnx export updates (#19692)

[zh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zha0q1 pushed a commit to branch v1.x
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git


The following commit(s) were added to refs/heads/v1.x by this push:
     new d06d705  [v1.x] More onnx export updates (#19692)
d06d705 is described below

commit d06d7054b47160f3308a0e6d86ad8edc86331b89
Author: Joe Evans <jo...@gmail.com>
AuthorDate: Thu Dec 17 21:45:53 2020 -0800

    [v1.x] More onnx export updates (#19692)
    
    * Comment out useless backend test for onnx.
    
    * Infer the output data type for each graph layer instead of assuming it's the same as input type.
    
    * Add unit test for cast onnx export operator.
    
    * Fix lint
    
    * Fix lint
    
    Co-authored-by: Joe Evans <jo...@amazon.com>
---
 ci/docker/runtime_functions.sh                   |  2 +-
 python/mxnet/contrib/onnx/mx2onnx/export_onnx.py | 19 ++++++++++++++-----
 tests/python-pytest/onnx/test_operators.py       |  9 +++++++++
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/ci/docker/runtime_functions.sh b/ci/docker/runtime_functions.sh
index cabe853..6849944 100755
--- a/ci/docker/runtime_functions.sh
+++ b/ci/docker/runtime_functions.sh
@@ -1277,7 +1277,7 @@ integrationtest_ubuntu_cpu_onnx() {
     export PYTHONPATH=./python/
     export MXNET_SUBGRAPH_VERBOSE=0
     export DMLC_LOG_STACK_TRACE_DEPTH=10
-    tests/python-pytest/onnx/backend_test.py
+    #tests/python-pytest/onnx/backend_test.py
     COV_ARG="--cov=./ --cov-report=xml --cov-append"
     pytest $COV_ARG --verbose tests/python-pytest/onnx/mxnet_export_test.py
     pytest $COV_ARG --verbose tests/python-pytest/onnx/test_models.py
diff --git a/python/mxnet/contrib/onnx/mx2onnx/export_onnx.py b/python/mxnet/contrib/onnx/mx2onnx/export_onnx.py
index bd25336..817a019 100644
--- a/python/mxnet/contrib/onnx/mx2onnx/export_onnx.py
+++ b/python/mxnet/contrib/onnx/mx2onnx/export_onnx.py
@@ -116,7 +116,7 @@ class MXNetGraph(object):
         return arg_params, aux_params
 
     @staticmethod
-    def get_outputs(sym, params, in_shape, in_label):
+    def get_outputs(sym, params, in_shape, in_label, in_type):
         """ Infer output shapes and return dictionary of output name to shape
 
         :param :class:`~mxnet.symbol.Symbol` sym: symbol to perform infer shape on
@@ -127,6 +127,7 @@ class MXNetGraph(object):
         :return: dictionary of output name to shape
         :rtype: dict of (str, tuple(int, ...))
         """
+        from onnx import mapping
         # remove any input listed in params from sym.list_inputs() and bind them to the input shapes provided
         # by user. Also remove in_label, which is the name of the label symbol that may have been used
         # as the label for loss during training.
@@ -146,8 +147,16 @@ class MXNetGraph(object):
                 out_names.append(name)
 
         assert len(out_shapes) == len(out_names)
+
+        # infer output types
+        args = {n: mapping.TENSOR_TYPE_TO_NP_TYPE[in_type] for n in sym.list_inputs()}
+        _, out_type, _ = sym.infer_type(**args)
+        out_types = [mapping.NP_TYPE_TO_TENSOR_TYPE[o(0).dtype] for o in out_type]
+
+        assert len(out_types) == len(out_names)
+
         # bind output shapes with output names
-        graph_outputs = {n: s for n, s in zip(out_names, out_shapes)}
+        graph_outputs = {n: {'shape': s, 'dtype': d} for n, s, d in zip(out_names, out_shapes, out_types)}
 
         return graph_outputs
 
@@ -210,7 +219,7 @@ class MXNetGraph(object):
         index_lookup = []
 
         # Determine output shape
-        graph_outputs = MXNetGraph.get_outputs(sym, params, in_shape, output_label)
+        graph_outputs = MXNetGraph.get_outputs(sym, params, in_shape, output_label, in_type)
 
         graph_input_idx = 0
         for idx, node in enumerate(mx_graph):
@@ -273,8 +282,8 @@ class MXNetGraph(object):
                                 onnx_processed_outputs.append(
                                     make_tensor_value_info(
                                         name=nodename,
-                                        elem_type=in_type,
-                                        shape=graph_outputs[nodename]
+                                        elem_type=graph_outputs[nodename]['dtype'],
+                                        shape=graph_outputs[nodename]['shape']
                                     )
                                 )
                                 if verbose:
diff --git a/tests/python-pytest/onnx/test_operators.py b/tests/python-pytest/onnx/test_operators.py
index 334de59..87554ed 100644
--- a/tests/python-pytest/onnx/test_operators.py
+++ b/tests/python-pytest/onnx/test_operators.py
@@ -293,3 +293,12 @@ def test_onnx_export_dropout(tmp_path, dtype, p):
     M = def_model('Dropout', p=p)
     x = mx.nd.array([[3,0.5,-0.5,2,7],[2,-0.4,7,3,0.2]], dtype=dtype)
     op_export_test('Dropout', M, [x], tmp_path)
+
+
+@pytest.mark.parametrize('src_dtype', ['float16', 'float32', 'float64'])
+@pytest.mark.parametrize('dst_dtype', ['bool', 'float16', 'float32', 'float64', 'int32', 'int64', 'int8', 'uint8'])
+@pytest.mark.parametrize('shape', [(2,3), (4,5,6)])
+def test_onnx_export_cast(tmp_path, src_dtype, dst_dtype, shape):
+    M = def_model('Cast', dtype=dst_dtype)
+    x = mx.nd.ones(shape, dtype=src_dtype)
+    op_export_test('Cast', M, [x], tmp_path)