Merge pull request #1221 from xadupre/squ

[WIP] Support operators Squeeze/Unsqueeze for opset 13 (v2)

Author: TomWildenhain-Microsoft

tests/test_optimizers.py

@@ -318,6 +318,7 @@ def test_transpose_with_identity(self):
         self.run_transpose_compare(["Z"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
                                    model_proto, remaining_transpose_num=1)
 
+    @check_opset_max_version(12, "Squeeze/Unsqueeze changed in opset 13")
     def test_transpose_with_squeeze1(self):
         # squeeze the first dim
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
@@ -335,6 +336,26 @@ def test_transpose_with_squeeze1(self):
                                                      model_proto, remaining_transpose_num=1)
         self.check_transpose_perm(model_after_opt, [1, 2, 0])
 
+    @check_opset_min_version(13, "Squeeze/Unsqueeze changed in opset 13")
+    def test_transpose_with_squeeze1_13(self):
+        # squeeze the first dim
+        node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
+        axes = self._make_onnx_const(np.array([0], dtype=np.int64), "axes")
+        node2 = helper.make_node("Squeeze", ["Y", "axes"], ["Z"], name="squeeze")
+
+        graph = helper.make_graph(
+            [node1, node2, axes],
+            "transpose_with_squeeze",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1, 3, 4, 5))],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (4, 5, 3))],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        model_after_opt = self.run_transpose_compare(["Z"], {"X": np.random.randn(1, 3, 4, 5).astype(np.float32)},
+                                                     model_proto, remaining_transpose_num=1)
+        self.check_transpose_perm(model_after_opt, [1, 2, 0])
+
+    @check_opset_max_version(12, "Squeeze/Unsqueeze changed in opset 13")
     def test_transpose_with_squeeze2(self):
         # squeeze the second dim
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
@@ -352,6 +373,26 @@ def test_transpose_with_squeeze2(self):
                                                      model_proto, remaining_transpose_num=1)
         self.check_transpose_perm(model_after_opt, [0, 2, 1])
 
+    @check_opset_min_version(13, "Squeeze/Unsqueeze changed in opset 13")
+    def test_transpose_with_squeeze2_13(self):
+        # squeeze the second dim
+        node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
+        axes = self._make_onnx_const(np.array([1], dtype=np.int64), "axes")
+        node2 = helper.make_node("Squeeze", ["Y", "axes"], ["Z"], name="squeeze")
+
+        graph = helper.make_graph(
+            [node1, node2, axes],
+            "transpose_with_squeeze",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (3, 4, 1, 5))],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (3, 5, 4))],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        model_after_opt = self.run_transpose_compare(["Z"], {"X": np.random.randn(3, 4, 1, 5).astype(np.float32)},
+                                                     model_proto, remaining_transpose_num=1)
+        self.check_transpose_perm(model_after_opt, [0, 2, 1])
+
+    @check_opset_max_version(12, "Squeeze/Unsqueeze changed in opset 13")
     def test_transpose_with_squeeze3(self):
         # squeeze the last dim
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
@@ -368,6 +409,25 @@ def test_transpose_with_squeeze3(self):
         self.run_transpose_compare(["Z"], {"X": np.random.randn(3, 1, 4, 5).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    @check_opset_min_version(13, "Squeeze/Unsqueeze changed in opset 13")
+    def test_transpose_with_squeeze3_13(self):
+        # squeeze the last dim
+        node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
+        axes = self._make_onnx_const(np.array([3], dtype=np.int64), "axes")
+        node2 = helper.make_node("Squeeze", ["Y", "axes"], ["Z"], name="squeeze")
+
+        graph = helper.make_graph(
+            [node1, node2, axes],
+            "transpose_with_squeeze",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (3, 1, 4, 5))],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (3, 4, 5))],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["Z"], {"X": np.random.randn(3, 1, 4, 5).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
+    @check_opset_max_version(12, "Squeeze/Unsqueeze changed in opset 13")
     def test_transpose_with_squeeze4(self):
         # squeeze the two dims
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
@@ -384,6 +444,24 @@ def test_transpose_with_squeeze4(self):
         self.run_transpose_compare(["Z"], {"X": np.random.randn(3, 1, 1, 5).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    @check_opset_min_version(13, "Squeeze/Unsqueeze changed in opset 13")
+    def test_transpose_with_squeeze4_13(self):
+        # squeeze the two dims
+        node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
+        axes = self._make_onnx_const(np.array([1, 3], dtype=np.int64), "axes")
+        node2 = helper.make_node("Squeeze", ["Y", "axes"], ["Z"], name="squeeze")
+
+        graph = helper.make_graph(
+            [node1, node2, axes],
+            "transpose_with_squeeze",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (3, 1, 1, 5))],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (3, 5))],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["Z"], {"X": np.random.randn(3, 1, 1, 5).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
     def test_transpose_with_loop(self):
         def _define_loop_graph(external_inputs):
             # external_inputs: external node which will be used by this graph
@@ -1090,6 +1168,7 @@ def test_const_fold_node_is_output(self):
         self.run_transpose_compare(["res"], {},
                                    model_proto, remaining_transpose_num=0)
 
+    @check_opset_max_version(12, "Squeeze/Unsqueeze changed in opset 13")
     def test_const_fold_unsqueeze_with_const(self):
         shape = (6, 6)
         const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,
@@ -1109,6 +1188,27 @@ def test_const_fold_unsqueeze_with_const(self):
         self.run_and_compare(["res"], {"X": np.random.randn(1).astype(np.float32)}, model_proto,
                              "Unsqueeze", 0)
 
+    @check_opset_min_version(13, "Squeeze/Unsqueeze changed in opset 13")
+    def test_const_fold_unsqueeze_with_const_13(self):
+        shape = (6, 6)
+        const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,
+                                          vals=np.random.randn(*shape).flatten().astype(np.float32))
+        node1 = helper.make_node("Constant", [], ["const"], value=const_tensor)
+        axes = self._make_onnx_const(np.array([0, 2, 3], dtype=np.int64), "axes")
+        node2 = helper.make_node("Unsqueeze", ["const", "axes"], ["value1"])
+        node3 = helper.make_node("Add", ["value1", "X"], ["res"])
+
+        graph = helper.make_graph(
+            [node1, node2, node3, axes],
+            "test_const_fold_unsqueeze_with_const",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (1,))],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, (1, 6, 1, 1, 6))],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_and_compare(["res"], {"X": np.random.randn(1).astype(np.float32)}, model_proto,
+                             "Unsqueeze", 0)
+
     def test_const_fold_cast_with_const(self):
         shape = (6, 6)
         const_tensor = helper.make_tensor(name='const_tensor', data_type=TensorProto.FLOAT, dims=shape,

tf2onnx/custom_opsets/string_ops.py

@@ -8,6 +8,7 @@
 from tf2onnx import constants, handler
 from tf2onnx.handler import tf_op
 from tf2onnx import utils
+from tf2onnx.graph_builder import GraphBuilder
 
 logger = logging.getLogger(__name__)
 
@@ -16,7 +17,7 @@
 @tf_op(["StringSplit", "StringSplitV2"], domain=constants.CONTRIB_OPS_DOMAIN)
 class StringOps:
     @classmethod
-    def version_1(cls, ctx, node, **kwargs):
+    def any_version(cls, opset, ctx, node, **kwargs):
         if node.type == "StringSplit":
             skip_empty = node.get_attr_value('skip_empty', True)
         else:
@@ -25,10 +26,21 @@ def version_1(cls, ctx, node, **kwargs):
         node.domain = constants.CONTRIB_OPS_DOMAIN
         for a in list(node.attr.keys()):
             del node.attr[a]
-        unsqueeze_node = ctx.make_node("Unsqueeze", [node.input[1]], attr={'axes': [0]})
+        unsqueeze_node = GraphBuilder(ctx).make_squeeze(
+            {'data': node.input[1], 'axes': [0]}, return_node=True)
+
         skip_empty_const = ctx.make_const(utils.make_name('skip_empty_const'), np.array([skip_empty], np.bool))
         ctx.replace_inputs(node, [node.input[0], unsqueeze_node.output[0], skip_empty_const.output[0]])
 
+    @classmethod
+    def version_1(cls, ctx, node, **kwargs):
+        cls.any_version(1, ctx, node, **kwargs)
+
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        cls.any_version(13, ctx, node, **kwargs)
+
+
 @tf_op("StringToHashBucketFast", domain=constants.CONTRIB_OPS_DOMAIN)
 class StringToHashBucketFast:
     @classmethod
@@ -59,7 +71,7 @@ def version_1(cls, ctx, node, **kwargs):
 @tf_op("StringJoin", domain=constants.CONTRIB_OPS_DOMAIN)
 class StringJoin:
     @classmethod
-    def version_1(cls, ctx, node, **kwargs):
+    def any_version(cls, opset, ctx, node, **kwargs):
         node.domain = constants.CONTRIB_OPS_DOMAIN
         separator = node.get_attr_value("separator")
         if separator is None:
@@ -76,11 +88,20 @@ def version_1(cls, ctx, node, **kwargs):
             if ctx.get_shape(inp) == [] and shape_node is not None:
                 expand_node = ctx.make_node("Expand", [inp, shape_node.output[0]])
                 inp = expand_node.output[0]
-            unsqueeze_node = ctx.make_node("Unsqueeze", [inp], attr={'axes': [0]})
+                unsqueeze_node = GraphBuilder(ctx).make_squeeze({'data': inp, 'axes': [0]})
             unsqueezes.append(unsqueeze_node.output[0])
         stack_node = ctx.make_node("Concat", unsqueezes, attr={'axis': 0})
         ctx.replace_inputs(node, [stack_node.output[0], separator_node.output[0], axis_node.output[0]])
 
+    @classmethod
+    def version_1(cls, ctx, node, **kwargs):
+        cls.any_version(1, ctx, node, **kwargs)
+
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        cls.any_version(13, ctx, node, **kwargs)
+
+
 @tf_op(["Equal", "NotEqual"], domain=constants.CONTRIB_OPS_DOMAIN)
 class StringEqual:
     @classmethod

tf2onnx/graph.py

@@ -1258,8 +1258,25 @@ def insert_new_node_on_input(self, node, op_type, input_name, name=None, domain=
                 break
         return new_node
 
-    def insert_new_node_on_output(self, op_type, output_name, name=None, inputs=None, domain=None, **kwargs):
+    def insert_node_on_output(self, node, output_name=None):
+        """
+        The inserted node takes the *output_name* as input and produces a
+        new output. The function goes through every node taking *output_name*
+        and replaces it by the new output name.
+        """
+        if output_name is None:
+            output_name = node.input[0]
+        new_output = node.output[0]
+
+        to_replace = [self.get_node_by_name(n) for n in self._output_to_consumers[output_name]]
+        to_replace = [n for n in to_replace if n != node]
+        self.replace_all_inputs(output_name, new_output, ops=to_replace)
+        return node
+
+    def insert_new_node_on_output(self, op_type, output_name=None, name=None, inputs=None, domain=None, **kwargs):
         """Create and insert a new node into the graph.
+        It then calls insert_node_on_output.
+
         Args:
             op_type: type for new operation
             output_name: the names of the outputs above us
@@ -1273,6 +1290,7 @@ def insert_new_node_on_output(self, op_type, output_name, name=None, inputs=None
                         type(output_name))
         utils.make_sure(isinstance(op_type, six.text_type), "op_type's type is not expected: %s",
                         type(op_type))
+        utils.make_sure(output_name is not None, "output_name cannot be None for op_type=%r.", op_type)
 
         if inputs is None:
             inputs = [output_name]
@@ -1281,11 +1299,7 @@ def insert_new_node_on_output(self, op_type, output_name, name=None, inputs=None
 
         new_output = port_name(name)
         new_node = self.make_node(op_type, inputs, attr=kwargs, outputs=[new_output], name=name, domain=domain)
-
-        to_replace = [self.get_node_by_name(n) for n in self._output_to_consumers[output_name]]
-        to_replace = [n for n in to_replace if n != new_node]
-        self.replace_all_inputs(output_name, new_output, ops=to_replace)
-        return new_node
+        return self.insert_node_on_output(new_node, output_name)
 
     def find_output_consumers(self, output_name):
         """Find all nodes consuming a given output."""

tf2onnx/graph_builder.py

@@ -24,7 +24,7 @@ def __init__(self, graph):
     def graph(self):
         return self._g
 
-    def make_slice(self, kwargs, name=None, shapes=None, dtypes=None):
+    def make_slice(self, kwargs, name=None, shapes=None, dtypes=None, return_node=False):
         """
         slice changes its schema at opset 10: it treats some attributes as dynamic input
         so this function has to process inputs according to graph's opset version
@@ -77,8 +77,11 @@ def make_slice(self, kwargs, name=None, shapes=None, dtypes=None):
                 if input_data != utils.ONNX_EMPTY_INPUT:
                     utils.make_sure(dtype == self.graph.get_dtype(input_data), "dtype should be same")
 
-        return self.graph.make_node(op_type="Slice", inputs=inputs, attr=attr, name=name,
-                                    outputs=outputs, shapes=shapes, dtypes=dtypes).output[0]
+        node = self.graph.make_node(op_type="Slice", inputs=inputs, attr=attr, name=name,
+                                    outputs=outputs, shapes=shapes, dtypes=dtypes)
+        if return_node:
+            return node
+        return node.output[0]
 
     def make_reduce_sum(self, kwargs, name=None, shapes=None, dtypes=None):
         """
@@ -116,7 +119,7 @@ def make_reduce_sum(self, kwargs, name=None, shapes=None, dtypes=None):
         return self.graph.make_node(op_type="ReduceSum", inputs=inputs, attr=attr, name=name,
                                     outputs=outputs, shapes=shapes, dtypes=dtypes).output[0]
 
-    def make_squeeze(self, kwargs, name=None, shapes=None, dtypes=None):
+    def make_squeeze(self, kwargs, name=None, shapes=None, dtypes=None, return_node=False, op_name_scope=None):
         """
         Squeeze changes its schema at opset 13: it treats axes as a dynamic input
         kwargs: key could be ["data", "axes"].
@@ -150,10 +153,14 @@ def make_squeeze(self, kwargs, name=None, shapes=None, dtypes=None):
         while inputs[-1] == utils.ONNX_EMPTY_INPUT:
             inputs = inputs[:-1]
 
-        return self.graph.make_node(op_type="Squeeze", inputs=inputs, attr=attr, name=name,
-                                    outputs=outputs, shapes=shapes, dtypes=dtypes).output[0]
+        node = self.graph.make_node(op_type="Squeeze", inputs=inputs, attr=attr, name=name,
+                                    outputs=outputs, shapes=shapes, dtypes=dtypes,
+                                    op_name_scope=op_name_scope)
+        if return_node:
+            return node
+        return node.output[0]
 
-    def make_unsqueeze(self, kwargs, name=None, shapes=None, dtypes=None):
+    def make_unsqueeze(self, kwargs, name=None, shapes=None, dtypes=None, return_node=False, op_name_scope=None):
         """
         Unsqueeze changes its schema at opset 13: it treats axes as a dynamic input
         kwargs: key could be ["data", "axes"].
@@ -187,8 +194,12 @@ def make_unsqueeze(self, kwargs, name=None, shapes=None, dtypes=None):
         while inputs[-1] == utils.ONNX_EMPTY_INPUT:
             inputs = inputs[:-1]
 
-        return self.graph.make_node(op_type="Unsqueeze", inputs=inputs, attr=attr, name=name,
-                                    outputs=outputs, shapes=shapes, dtypes=dtypes).output[0]
+        node = self.graph.make_node(op_type="Unsqueeze", inputs=inputs, attr=attr, name=name,
+                                    outputs=outputs, shapes=shapes, dtypes=dtypes,
+                                    op_name_scope=op_name_scope)
+        if return_node:
+            return node
+        return node.output[0]
 
     def convert_to_input(self, tensor, const_name, is_optional=False, dtype=None):
         """in ONNX, input shold come from node, so it must be a string"""

tf2onnx/onnx_opset/controlflow.py

@@ -19,6 +19,7 @@
 from tf2onnx import utils
 from tf2onnx.handler import tf_op
 from tf2onnx.tf_loader import find_function
+from tf2onnx.graph_builder import GraphBuilder
 
 
 logger = logging.getLogger(__name__)
@@ -289,6 +290,20 @@ def version_7(cls, ctx, node, **kwargs):
         ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[1], name=node.child_name(), axes=[0])
         ctx.insert_new_node_on_output("Squeeze", node.output[0], name=node.child_name(), axes=[0])
 
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        ctx.ta_reads.append(node.input[0])
+        node.type = "Gather"
+        ctx.replace_inputs(node, [node.input[0], node.input[1]])
+
+        g = GraphBuilder(ctx)
+
+        usq_node = g.make_unsqueeze({"axes": [0], 'name': node.child_name(), 'data': node.input[1]}, return_node=True)
+        ctx.insert_node_on_output(usq_node)
+
+        sq_node = g.make_squeeze({"axes": [0], 'name': node.child_name(), 'data': node.output[0]}, return_node=True)
+        ctx.insert_node_on_output(sq_node)
+
 
 @tf_op(["TensorListLength"])
 class TensorListLength:
@@ -607,7 +622,7 @@ def inline_subgraph(parent, g, scope, binding):
         parent.set_dtype(name, g.get_dtype(name))
         parent.set_shape(name, g.get_shape(name))
 
-    return  g.outputs
+    return g.outputs
 
 
 def parameter_binding(g, inputs, state_vars=None):