Merge pull request #1216 from onnx/tom/SqueezeUnsqueeze

Tom/squeeze unsqueeze

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1044,12 +1044,31 @@ def func(x1, x2):
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2})
 
     def test_sequeeze_no_axis_specified(self):
-        x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 2, 1))
+        x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 1, 2, 1, 1))
+        def func(x):
+            x_ = tf.squeeze(x)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
+    def test_sequeeze_no_axis(self):
+        x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 2))
         def func(x):
             x_ = tf.squeeze(x)
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(11, "Pad")
+    def test_sequeeze_no_axis_specified_unknown_rank(self):
+        x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
+        y_val = np.array([2, 1, 2, 1, 1], dtype=np.int64)
+        z_val = np.zeros((1, 2), dtype=np.int64)
+        def func(x, y, z):
+            y_ = tf.pad(y, z)
+            x_ = tf.reshape(x, y_)
+            x_ = tf.squeeze(x_)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val, _INPUT2: z_val})
+
     def test_sequeeze_positive_axis(self):
         x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 2, 1))
         def func(x):
@@ -1071,6 +1090,18 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(11, "Squeeze")
+    def test_sequeeze_mixed_axis_unknown_rank(self):
+        x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
+        y_val = np.array([2, 1, 2, 1, 1], dtype=np.int64)
+        z_val = np.zeros((1, 2), dtype=np.int64)
+        def func(x, y, z):
+            y_ = tf.pad(y, z)
+            x_ = tf.reshape(x, y_)
+            x_ = tf.squeeze(x_, [1, -1])
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val, _INPUT2: z_val})
+
     def test_transpose(self):
         x_val = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], dtype=np.float32).reshape((2, 3))
         def func(x):

tf2onnx/graph_builder.py

@@ -116,6 +116,80 @@ 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):
+        """
+        Squeeze changes its schema at opset 13: it treats axes as a dynamic input
+        kwargs: key could be ["data", "axes"].
+        """
+        outputs = kwargs.pop("outputs", None)
+
+        if self.graph.opset < 13:
+            data = kwargs.pop("data")
+            axes = self.convert_to_attribute(kwargs.pop("axes", None), is_optional=True)
+            attr = {"axes": axes}
+            inputs = [data]
+        else:
+            data = kwargs.pop("data")
+            axes = self.convert_to_input(kwargs.pop("axes", None), "const_axes", is_optional=True, dtype=np.int64)
+            attr = {}
+            inputs = [data, axes]
+
+        if kwargs:
+            logger.warning("kwargs contains un-used key")
+
+        new_attr = {}
+        for key, val in attr.items():
+            if val is not None:
+                new_attr[key] = val
+        attr = new_attr
+
+        for ind, val in enumerate(inputs):
+            if val is None:
+                inputs[ind] = utils.ONNX_EMPTY_INPUT  # empty string means no connection in ONNX
+        # remove tailing ""
+        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]
+
+    def make_unsqueeze(self, kwargs, name=None, shapes=None, dtypes=None):
+        """
+        Unsqueeze changes its schema at opset 13: it treats axes as a dynamic input
+        kwargs: key could be ["data", "axes"].
+        """
+        outputs = kwargs.pop("outputs", None)
+
+        if self.graph.opset < 13:
+            data = kwargs.pop("data")
+            axes = self.convert_to_attribute(kwargs.pop("axes", None), is_optional=True)
+            attr = {"axes": axes}
+            inputs = [data]
+        else:
+            data = kwargs.pop("data")
+            axes = self.convert_to_input(kwargs.pop("axes", None), "const_axes", is_optional=True, dtype=np.int64)
+            attr = {}
+            inputs = [data, axes]
+
+        if kwargs:
+            logger.warning("kwargs contains un-used key")
+
+        new_attr = {}
+        for key, val in attr.items():
+            if val is not None:
+                new_attr[key] = val
+        attr = new_attr
+
+        for ind, val in enumerate(inputs):
+            if val is None:
+                inputs[ind] = utils.ONNX_EMPTY_INPUT  # empty string means no connection in ONNX
+        # remove tailing ""
+        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]
+
     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"""
         if is_optional and tensor is None:

tf2onnx/onnx_opset/tensor.py

@@ -183,28 +183,25 @@ class Squeeze:
     def version_1(cls, ctx, node, **kwargs):
         # T output = Squeeze(T input, @list(int) squeeze_dims)
         # T squeezed = Squeeze(T data, @AttrType.INTS axes), axes are list of positive integers.
-        axis = node.get_attr("axis")
-        if not axis:
-            axis = node.get_attr("squeeze_dims")
-            if axis:
-                del node.attr["squeeze_dims"]
+        axes = node.get_attr_value("squeeze_dims")
+        if axes is None:
+            axes = []
         else:
-            del node.attr["axis"]
+            del node.attr["squeeze_dims"]
 
-        if axis and axis.ints:
-            axis = axis.ints
-            neg_axis = any([val < 0 for val in axis])
-            if neg_axis:
+        # TF uses empty axes to indicate that all 1 dims should be squeezed
+        if len(axes) > 0:
+            neg_axis = any([val < 0 for val in axes])
+            if neg_axis and ctx.opset < 11:
                 shape = ctx.get_shape(node.input[0])
-                utils.make_sure(shape is not None, "squeeze input shape cannot be None")
+                utils.make_sure(shape is not None, "squeeze with negative axes and unknown rank requires opset >= 11")
                 shape_len = len(shape)
-                axis = [a + shape_len if a < 0 else a for a in axis]
-        else:
-            shape = ctx.get_shape(node.input[0])
-            utils.make_sure(shape is not None, "squeeze input shape cannot be None")
-            axis = [i for i, j in enumerate(shape) if j == 1]
-            if not axis: axis = [0]
-        node.set_attr("axes", axis)
+                axes = [a + shape_len if a < 0 else a for a in axes]
+            if ctx.opset < 13:
+                node.set_attr("axes", axes)
+            else:
+                axes_const = ctx.make_const(utils.make_name("axes_const"), np.array(axes, dtype=np.int64))
+                ctx.replace_inputs(node, [node.input[0], axes_const.output[0]])
 
     @classmethod
     def version_11(cls, ctx, node, **kwargs):