Merge pull request #346 from zhijxu-MS/push_branch

optimize upsample's conversion logic

Author: nbcsm

tests/backend_test_base.py

@@ -150,6 +150,8 @@ def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port
             if check_shape:
                 self.assertEqual(expected_val.shape, actual_val.shape)
 
+        return g
+
     def save_onnx_model(self, model_proto, feed_dict, postfix=""):
         target_path = utils.save_onnx_model(self.test_data_directory, self._testMethodName + postfix, feed_dict,
                                             model_proto, include_test_data=self.config.is_debug_mode,

tests/common.py

@@ -7,6 +7,7 @@
 import os
 import sys
 import unittest
+from collections import defaultdict
 
 from distutils.version import LooseVersion
 from tf2onnx import utils
@@ -15,7 +16,7 @@
 __all__ = ["TestConfig", "get_test_config", "unittest_main",
            "check_tf_min_version", "skip_tf_versions",
            "check_opset_min_version", "check_target", "skip_onnxruntime_backend", "skip_caffe2_backend",
-           "check_onnxruntime_incompatibility"]
+           "check_onnxruntime_incompatibility", "validate_const_node", "group_nodes_by_type"]
 
 
 # pylint: disable=missing-docstring
@@ -209,3 +210,17 @@ def check_onnxruntime_incompatibility(op):
 
     reason = "{} is not supported by onnxruntime before opset {}".format(op, support_since[op])
     return unittest.skipIf(True, reason)
+
+
+def validate_const_node(node, expected_val):
+    if node.is_const():
+        node_val = node.get_tensor_value()
+        return node_val == expected_val
+    return False
+
+
+def group_nodes_by_type(graph):
+    res = defaultdict(list)
+    for node in graph.get_nodes():
+        res[node.type].append(node)
+    return res

tests/test_backend.py

@@ -92,7 +92,7 @@ class BackendTests(Tf2OnnxBackendTestBase):
     def _run_test_case(self, output_names_with_port, feed_dict, **kwargs):
         kwargs["convert_var_to_const"] = False
         kwargs["constant_fold"] = False
-        self.run_test_case(feed_dict, [], output_names_with_port, **kwargs)
+        return self.run_test_case(feed_dict, [], output_names_with_port, **kwargs)
 
     def _test_expand_dims(self, idx):
         tf.reset_default_graph()
@@ -1276,7 +1276,10 @@ def test_resize_nearest_neighbor(self):
         x_new_size_ = tf.constant(x_new_size)
         x_ = tf.image.resize_nearest_neighbor(x, x_new_size_)
         _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        graph = self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        if self.config.opset >= 9:
+            scale_node = group_nodes_by_type(graph)["Upsample"][0].inputs[1]
+            self.assertTrue(validate_const_node(scale_node, [1.0, 1.0, 2.0, 2.0]))
 
     @check_opset_min_version(9, "resize_nearest_neighbor")
     def test_resize_nearest_neighbor_with_non_const(self):
@@ -1301,7 +1304,10 @@ def test_resize_bilinear(self):
         x_new_size_ = tf.constant(x_new_size)
         x_ = tf.image.resize_bilinear(x, x_new_size_)
         _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        graph = self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        if self.config.opset >= 9:
+            scale_node = group_nodes_by_type(graph)["Upsample"][0].inputs[1]
+            self.assertTrue(validate_const_node(scale_node, [1.0, 1.0, 2.0, 2.0]))
 
     @check_opset_min_version(9, "resize_bilinear")
     def test_resize_bilinear_with_non_const(self):

tf2onnx/graph.py

@@ -373,14 +373,19 @@ def is_target(self, *names):
         """Return True if target platform contains any name."""
         return any(name in self._target for name in names)
 
-    def make_const(self, name, np_val, skip_conversion=False):
+    def make_const(self, name, np_val, skip_conversion=False, raw=True):
         """Make a new constant in the graph.
         Args:
             name: const node name, must be unique.
             np_val: value of type numpy ndarray.
             skip_conversion: bool, indicate whether this created node would be mapped during conversion.
+            raw: whether to store data at field of raw_data or the specific field according to its dtype
         """
-        onnx_tensor = numpy_helper.from_array(np_val, name)
+        if raw:
+            onnx_tensor = numpy_helper.from_array(np_val, name)
+        else:
+            onnx_tensor = helper.make_tensor(name, utils.map_numpy_to_onnx_dtype(np_val.dtype),
+                                             np_val.shape, np_val, raw=False)
         node = self.make_node("Const", [], outputs=[name], name=name, attr={"value": onnx_tensor},
                               skip_conversion=skip_conversion)
         self.set_shape(name, np_val.shape)
@@ -732,8 +737,12 @@ def make_graph(self, doc, graph_name="tf2onnx"):
         # create initializers for constant nodes
         const_ops = [op for op in const_ops if op not in placeholder_default_const_ops]
         for op in const_ops:
-            const_val = op.get_tensor_value(as_list=False)
-            tensor = numpy_helper.from_array(const_val, op.output[0])
+            # not to use numpy_helper.from_array to create a new tensor
+            # because sometimes onnx will have a bug that only check the tensor data in specific field
+            # such as at upsample it only checks the float_data field.
+            t = op.get_attr("value")
+            tensor = helper.get_attribute_value(t)
+            tensor.name = op.output[0]
             initializers.append(tensor)
 
         # create input_tensor_values

tf2onnx/tfonnx.py

@@ -1056,18 +1056,32 @@ def upsample_op9(ctx, node, name, args):
     # float32 out = ResizeBilinear/ResizeNearestNeighbor(T images, int size)
     # https://www.tensorflow.org/api_docs/python/tf/image/resize_nearest_neighbor
     # wants the input to be NHWC - adjust target_shape to this.
-    ori_shape = ctx.make_node("Shape", [node.input[0]])
-    ori_shape_hw = ctx.make_node("Slice", ori_shape.output, {"axes": [0], "starts": [1], "ends": [3]})
-    ori_shape_hw_float = ctx.make_node("Cast", ori_shape_hw.output, attr={"to": onnx_pb.TensorProto.FLOAT})
 
-    target_hw = node.inputs[1]
-    target_hw_float = ctx.make_node("Cast", target_hw.output, attr={"to": onnx_pb.TensorProto.FLOAT})
+    # first create "scales" info for onnx upsample
+    # if shape of input and output known then  "scale" is calculated statically and set as a const node
+    shape = ctx.get_shape(node.input[0])
+    if shape and shape[2] != -1 and shape[1] != -1 and node.inputs[1].is_const():
+        target_shape = node.inputs[1].get_tensor_value()
+        n, h, w, c = shape
+        nh, nw = target_shape
+        # scales is nchw
+        # the reason not storing data at raw field is because of the bug: https://github.com/onnx/onnx/issues/1852
+        scale_val = np.array([1.0, 1.0, float(nh) / h, float(nw) / w]).astype(np.float32)
+        scales = ctx.make_const(utils.make_name("scales"), scale_val, raw=False)
+    else:
+        ori_shape = ctx.make_node("Shape", [node.input[0]])
+        ori_shape_hw = ctx.make_node("Slice", ori_shape.output, {"axes": [0], "starts": [1], "ends": [3]})
+        ori_shape_hw_float = ctx.make_node("Cast", ori_shape_hw.output, attr={"to": onnx_pb.TensorProto.FLOAT})
 
-    scales_hw = ctx.make_node("Div", [target_hw_float.output[0], ori_shape_hw_float.output[0]])
+        target_hw = node.inputs[1]
+        target_hw_float = ctx.make_node("Cast", target_hw.output, attr={"to": onnx_pb.TensorProto.FLOAT})
 
-    const_one_array = ctx.make_const(utils.make_name("one"), np.array([1.0, 1.0]).astype(np.float32))
-    # scaler is nchw
-    scales = ctx.make_node("Concat", [const_one_array.output[0], scales_hw.output[0]], {"axis": 0})
+        scales_hw = ctx.make_node("Div", [target_hw_float.output[0], ori_shape_hw_float.output[0]])
+
+        const_one_array = ctx.make_const(utils.make_name("one"), np.array([1.0, 1.0]).astype(np.float32))
+        # scales is nchw
+        scales = ctx.make_node("Concat", [const_one_array.output[0], scales_hw.output[0]], {"axis": 0})
+    # because onnxruntime only supports to scale the last two dims so transpose is inserted
     input_nchw = ctx.make_node("Transpose", [node.input[0]], {"perm": [0, 3, 1, 2]})
     upsample = ctx.make_node("Upsample", [input_nchw.output[0], scales.output[0]], attr={"mode": args[0]})