Merge pull request #481 from zhijxu-MS/push_branch

support resize, nonmaxsuppression in opset 10; and fix some bugs.

Author: nbcsm

tests/test_backend.py

@@ -1529,7 +1529,10 @@ def test_resize_nearest_neighbor(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         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]
+            # in opset 10, upsample is removed and resize is defined.
+            node_statistic = group_nodes_by_type(graph)
+            mapped_node = (node_statistic.get("Upsample") or node_statistic.get("Resize"))[0]
+            scale_node = mapped_node.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")
@@ -1557,7 +1560,10 @@ def test_resize_bilinear(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         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]
+            # in opset 10, upsample is removed and resize is defined.
+            node_statistic = group_nodes_by_type(graph)
+            mapped_node = (node_statistic.get("Upsample") or node_statistic.get("Resize"))[0]
+            scale_node = mapped_node.inputs[1]
             self.assertTrue(validate_const_node(scale_node, [1.0, 1.0, 2.0, 2.0]))
 
     @check_opset_min_version(9, "resize_bilinear")
@@ -1573,6 +1579,35 @@ def test_resize_bilinear_with_non_const(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: x_new_size})
 
+    @check_opset_min_version(10, "resize scale can less than 1")
+    def test_resize_bilinear_with_non_const2(self):
+        # scales has an element larger than 1 and also has an element less that 1
+        x_shape = [3, 100, 8, 5]
+        x_val = np.arange(1, 1 + np.prod(x_shape), dtype=np.float32).reshape(x_shape)
+        x = tf.placeholder(tf.float32, x_shape, name=_TFINPUT)
+
+        x_new_size = np.array([20, 16]).astype(np.int32)
+        x_new_size_ = tf.placeholder(shape=[None], dtype=tf.int32, name=_TFINPUT1)
+
+        x_ = tf.image.resize_bilinear(x, x_new_size_)
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: x_new_size})
+
+    @check_opset_min_version(10, "resize scale can less than 1")
+    def test_resize_nearest_neighbor2(self):
+        x_shape = [1, 300, 20, 2]
+        x_new_size = [30, 40]
+        x_val = np.arange(1, 1 + np.prod(x_shape)).astype("float32").reshape(x_shape)
+        x = tf.placeholder(tf.float32, x_shape, name=_TFINPUT)
+        x_new_size_ = tf.constant(x_new_size)
+        x_ = tf.image.resize_nearest_neighbor(x, x_new_size_)
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        graph = self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        node_statistic = group_nodes_by_type(graph)
+        mapped_node = node_statistic.get("Resize")[0]
+        scale_node = mapped_node.inputs[1]
+        self.assertTrue(validate_const_node(scale_node, [1.0, 1.0, 0.1, 2.0]))
+
     @check_opset_min_version(9, "fill")
     def test_fill_float32(self):
         x_shape = [1, 15, 20, 2]
@@ -2117,5 +2152,19 @@ def test_isinf(self):
                 self._run_test_case([_OUTPUT], {_INPUT: x_val})
                 tf.reset_default_graph()
 
+    @check_opset_min_version(10, "NonMaxSuppression")
+    def test_non_max_suppression(self):
+        box_num = 10
+        boxes_val = np.random.random_sample([box_num, 4]).astype(np.float32)
+        scores_val = np.random.random_sample([box_num]).astype(np.float32)
+        boxes = tf.placeholder(tf.float32, shape=[None, 4], name=_TFINPUT)
+        scores = tf.placeholder(tf.float32, shape=[None], name=_TFINPUT1)
+        res1 = tf.image.non_max_suppression(boxes, scores, max_output_size=int(box_num / 2))
+        res2 = tf.image.non_max_suppression(boxes, scores, max_output_size=0)
+        _ = tf.identity(res1, name=_TFOUTPUT)
+        _ = tf.identity(res2, name=_TFOUTPUT1)
+        self._run_test_case([_OUTPUT, _OUTPUT1], {_INPUT: boxes_val, _INPUT1: scores_val})
+
+
 if __name__ == '__main__':
     unittest_main()

tf2onnx/handler.py

@@ -15,6 +15,7 @@
 
 # pylint: disable=unused-argument,missing-docstring,invalid-name
 
+
 class tf_op:
     """Class to implement the decorator to register handlers that map tf to onnx."""
 

tf2onnx/onnx_opset/nn.py

@@ -464,14 +464,15 @@ def version_7(cls, ctx, node, **kwargs):
 
     @classmethod
     def version_9(cls, ctx, node, **kwargs):
-        cls._convert_since_9(ctx, node, **kwargs)
+        cls._convert_since_9(ctx, node, op_type="Upsample")
 
     @classmethod
     def version_10(cls, ctx, node, **kwargs):
-        cls._convert_since_9(ctx, node, **kwargs)
+        cls._convert_since_9(ctx, node, op_type="Resize")
 
     @classmethod
-    def _convert_since_9(cls, ctx, node, **kwargs):
+    def _convert_since_9(cls, ctx, node, op_type):
+
         # 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.
@@ -505,7 +506,7 @@ def _convert_since_9(cls, ctx, node, **kwargs):
             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": mode})
+        upsample = ctx.make_node(op_type, [input_nchw.output[0], scales.output[0]], attr={"mode": mode})
 
         shapes = node.output_shapes
         dtypes = node.output_dtypes

tf2onnx/onnx_opset/tensor.py

@@ -24,6 +24,7 @@
 
 # pylint: disable=unused-argument,missing-docstring,unused-variable
 
+
 def _convert_shapenode_to_int64(ctx, node, input_number):
     """cast int32 shape into int64 shape."""
     name = node.input[input_number]
@@ -55,6 +56,7 @@ def _wrap_concat_with_cast(ctx, node):
             ctx.set_dtype(output_cast.output[0], dtype)
             ctx.copy_shape(output_name, output_cast.output[0])
 
+
 @tf_op(["Size", "Flatten", "Dropout"])
 class DirectOp:
     @classmethod
@@ -296,8 +298,6 @@ def version_4(cls, ctx, node, **kwargs):
         ctx.remove_input(node, node.input[2])
         node.set_attr("axis", axis)
 
-INT64_MAX = np.iinfo(np.int64).max
-
 
 def _make_gathernd_inner_loop(ctx, params, index, dtype):
     """create the inner loop for GatherNd."""
@@ -343,7 +343,7 @@ def make_gathernd(ctx, params, indices, output, scope_name, t_params, shapes, dt
     # reshape indices into [sum(indices[:-1]), indices[-1]]
     indices_shape = ctx.make_node("Shape", [indices], dtypes=[TensorProto.INT64])
     indices_size = ctx.make_node("Size", [indices])
-    attr = {"axes": [0], "ends": [INT64_MAX], "starts": [-1]}
+    attr = {"axes": [0], "ends": [utils.get_max_value(np.int64)], "starts": [-1]}
     inputs_map = {"data": indices_shape.output[0], **attr}
     inner_shape = GraphBuilder(ctx).make_slice(inputs_map, dtypes=[TensorProto.INT64])
     outter_shape = ctx.make_node("Div",
@@ -401,7 +401,7 @@ def make_gathernd(ctx, params, indices, output, scope_name, t_params, shapes, dt
                                       [inner_loop_shape.output[0], one_const.output[0]],
                                       attr={"axis": 0},
                                       dtypes=[TensorProto.INT64])
-    attr = {"axes": [0], "ends": [INT64_MAX], "starts": [1]}
+    attr = {"axes": [0], "ends": [utils.get_max_value(np.int64)], "starts": [1]}
     inputs_map = {"data": inner_loop_shape_.output[0], **attr}
     output_inner_shape = GraphBuilder(ctx).make_slice(inputs_map, dtypes=[TensorProto.INT64])
     attr = {"axes": [0], "ends": [-1], "starts": [0]}
@@ -932,3 +932,29 @@ def version_10(cls, ctx, node, **kwargs):
         utils.make_sure(node_dtype, "Dtype of {} is None".format(node.name))
         if node_dtype not in [onnx_pb.TensorProto.FLOAT, onnx_pb.TensorProto.DOUBLE]:
             raise ValueError("dtype " + str(node_dtype) + " is not supported in onnx for now")
+
+
+@tf_op(["NonMaxSuppressionV2", "NonMaxSuppressionV3"], onnx_op="NonMaxSuppression")
+class NonMaxSuppression:
+    @classmethod
+    def version_10(cls, ctx, node, **kwargs):
+        # int32 = NonMaxSuppressionV2(T boxes, T scores, int32 max_output_size, T iou_threshold, T score_threshold)
+        # int64 = NonMaxSuppression(T boxes, T scores, int64 max_output_size, T iou_threshold, T score_threshold),
+        # T means float32 here, the last 3 params are optional
+        # tf boxes is 2D ([boxes_num, 4]) while onnx is 3D ([num_batches, boxes_num, 4])
+        # tf scores is 1D ([boxes_num])while onnx is 2D ([num_batches, num_classes, boxes_num])
+        # onnx output is [num_selected_boxes, 3], the meaning of last dim is [batch_index, class_index, box_index]
+        # while tf's output is [num_selected_boxes]
+        ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[0], axes=[0])
+        ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[1], axes=[0, 1])
+        ctx.insert_new_node_on_input(node, "Cast", node.input[2], to=onnx_pb.TensorProto.INT64)
+        # replace original node with nonmaxsurppress + slice + squeeze + cast
+        dtypes = [ctx.get_dtype(node.output[0])]
+        shapes = [ctx.get_shape(node.output[0])]
+        ctx.remove_node(node.name)
+        new_nonmaxsurppress = ctx.make_node(node.type, node.input).output[0]
+        slice_op = GraphBuilder(ctx).make_slice({"data": new_nonmaxsurppress,
+                                                 "axes": [1], "ends": [3], "starts": [2]})
+        squeeze_op = ctx.make_node("Squeeze", [slice_op], attr={"axes": [1]})
+        ctx.make_node("Cast", inputs=squeeze_op.output, attr={"to": onnx_pb.TensorProto.INT32},
+                      name=node.name, outputs=node.output, dtypes=dtypes, shapes=shapes)

tf2onnx/optimizer/transpose_optimizer.py

@@ -479,7 +479,7 @@ def _shape_handler(self, trans, node):
         self._g.remove_node(trans.name)
         self._g.remove_node(node.name)
         shape_node = self._g.make_node("Shape", [trans.input[0]])
-        const_node = self._g.make_const("Const", np.array(trans.get_attr("perm").ints))
+        const_node = self._g.make_const(utils.make_name("Const"), np.array(trans.get_attr("perm").ints))
         gather_node = self._g.make_node("Gather", [shape_node.output[0], const_node.output[0]], outputs=node.output)
         self._g.set_shape(gather_node.output[0], output_shape)
         self._g.set_dtype(gather_node.output[0], output_dtype)

tf2onnx/tfonnx.py

@@ -471,9 +471,7 @@ def rewrite_incomplete_type_support_rs5(g, ops):
 def rewrite_incomplete_type_support_rs6(g, ops):
     impacted_ops = [
         "Div",
-        "Greater",
         "IsNaN",
-        "Less",
         "Max",
         "Min",
         "ReduceSum",