Bug fixes for opset 13 support

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_optimizers.py

@@ -11,7 +11,7 @@
 import numpy as np
 from onnx import helper, TensorProto, OperatorSetIdProto
 from backend_test_base import Tf2OnnxBackendTestBase
-from common import unittest_main, group_nodes_by_type, check_opset_min_version, check_opset_max_version
+from common import unittest_main, group_nodes_by_type, check_opset_min_version, check_opset_max_version, get_test_config
 from tf2onnx import utils, constants
 from tf2onnx.graph import GraphUtil
 
@@ -471,12 +471,23 @@ def _define_loop_graph(external_inputs):
             node1 = helper.make_node("Gather", [external_inputs[0], "loop_iter_num"], ["Y0"])
             node2 = helper.make_node("Transpose", ["Y0"], ["Z0"], perm=[0, 2, 3, 1])
             # graph output
-            node3 = helper.make_node("Squeeze", ["Z0"], ["scan_output"], axes=[0])
+            if get_test_config().opset <= 12:
+                node3 = helper.make_node("Squeeze", ["Z0"], ["scan_output"], axes=[0])
+                const_node = None
+            else:
+                const_tensor = helper.make_tensor(name='const', data_type=TensorProto.INT64, dims=[1],
+                                                  vals=np.array([0], dtype=np.int64))
+                const_node = helper.make_node("Constant", [], ["axes_const"], value=const_tensor, name="const")
+                node3 = helper.make_node("Squeeze", ["Z0", "axes_const"], ["scan_output"])
             node4 = helper.make_node("Identity", ["loop_condition"], ["loop_cond_output"])
             node5 = helper.make_node("Identity", ["loop_condition"], ["loop_carried_output"])
 
+            nodes = [node1, node2, node3, node4, node5]
+            if const_node is not None:
+                nodes.append(const_node)
+
             graph = helper.make_graph(
-                [node1, node2, node3, node4, node5],
+                nodes,
                 "loop_subgraph",
                 [helper.make_tensor_value_info("loop_iter_num", TensorProto.INT64, (1,)),  # iteration_num
                  helper.make_tensor_value_info("loop_condition", TensorProto.BOOL, ()),  # condition
@@ -973,9 +984,9 @@ def test_duplicated_duplicated_input(self):
 
     def test_duplicated_duplicated_attributes(self):
         # same attr or not
-        node0 = helper.make_node('ReduceSum', inputs=["X"], outputs=["value0"], axes=[0], keepdims=0)
-        node1 = helper.make_node('ReduceSum', inputs=["X"], outputs=["value1"], axes=[0], keepdims=0)
-        node2 = helper.make_node('ReduceSum', inputs=["X"], outputs=["value2"], axes=[1], keepdims=0)
+        node0 = helper.make_node('ReduceMin', inputs=["X"], outputs=["value0"], axes=[0], keepdims=0)
+        node1 = helper.make_node('ReduceMin', inputs=["X"], outputs=["value1"], axes=[0], keepdims=0)
+        node2 = helper.make_node('ReduceMin', inputs=["X"], outputs=["value2"], axes=[1], keepdims=0)
         node3 = helper.make_node('Add', inputs=["value0", "value1"], outputs=["value3"])
         node4 = helper.make_node("Mul", ["value2", "value3"], ["OUT"])
 
@@ -988,7 +999,7 @@ def test_duplicated_duplicated_attributes(self):
 
         model_proto = self.make_model(graph, producer_name="onnx-tests")
         self.run_merge_duplicated_nodes_compare(["OUT"], {"X": np.random.randn(5, 5).astype(np.float32)}, model_proto,
-                                                op_type="ReduceSum", remaining_op_num=2)
+                                                op_type="ReduceMin", remaining_op_num=2)
 
     def _check_initializer_num(self, graph_proto, num):
         return num == len(graph_proto.initializer)

tf2onnx/graph_builder.py

@@ -55,8 +55,7 @@ def make_slice(self, kwargs, name=None, shapes=None, dtypes=None, return_node=Fa
             inputs = [data, starts, ends, axes, steps]
 
         # pro-process inputs and attr
-        if kwargs:
-            logger.warning("kwargs contains un-used key")
+        utils.make_sure(not kwargs, "kwargs contains un-used key")
 
         new_attr = {}
         for key, val in attr.items():
@@ -107,8 +106,7 @@ def make_reduce_sum(self, kwargs, name=None, shapes=None, dtypes=None):
             attr = {"keepdims": keepdims, "noop_with_empty_axes": noop_with_empty_axes}
             inputs = [data, axes]
 
-        if kwargs:
-            logger.warning("kwargs contains un-used key")
+        utils.make_sure(not kwargs, "kwargs contains un-used key")
 
         new_attr = {}
         for key, val in attr.items():
@@ -137,8 +135,7 @@ def make_squeeze(self, kwargs, name=None, shapes=None, dtypes=None, return_node=
             attr = {}
             inputs = [data, axes]
 
-        if kwargs:
-            logger.warning("kwargs contains un-used key")
+        utils.make_sure(not kwargs, "kwargs contains un-used key")
 
         new_attr = {}
         for key, val in attr.items():
@@ -178,8 +175,7 @@ def make_unsqueeze(self, kwargs, name=None, shapes=None, dtypes=None, return_nod
             attr = {}
             inputs = [data, axes]
 
-        if kwargs:
-            logger.warning("kwargs contains un-used key")
+        utils.make_sure(not kwargs, "kwargs contains un-used key")
 
         new_attr = {}
         for key, val in attr.items():

tf2onnx/onnx_opset/controlflow.py

@@ -298,10 +298,10 @@ def version_13(cls, ctx, node, **kwargs):
 
         g = GraphBuilder(ctx)
 
-        usq_node = g.make_unsqueeze({"axes": [0], 'name': node.child_name(), 'data': node.input[1]}, return_node=True)
+        usq_node = g.make_unsqueeze({"axes": [0], 'data': node.input[1]}, name=node.child_name(), 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)
+        sq_node = g.make_squeeze({"axes": [0], 'data': node.output[0]}, name=node.child_name(), return_node=True)
         ctx.insert_node_on_output(sq_node)
 
 

tf2onnx/onnx_opset/tensor.py

@@ -805,8 +805,10 @@ def version_1(cls, ctx, node, **kwargs):
             # insert_new_node_on_output(self, op_type, output_name=None, name=None, inputs=None, domain=None, **kwargs)
             # ctx.insert_new_node_on_output("Squeeze", node.output[0], name)
             name = utils.make_name(node.name)
-            squeeze_node = ctx.insert_new_node_on_output("Squeeze", node.output[0], name)
-            squeeze_node.set_attr("axes", needs_squeeze)
+            squeeze_node = GraphBuilder(ctx).make_squeeze(
+                {"axes": needs_squeeze, 'data': node.output[0]}, name=name, return_node=True)
+            ctx.insert_node_on_output(squeeze_node)
+
             nodes.append(squeeze_node)
             input_dtype = ctx.get_dtype(node.output[0])
             ctx.set_dtype(squeeze_node.output[0], input_dtype)
@@ -1023,8 +1025,9 @@ def any_version_after10(cls, opset, ctx, node, **kwargs):
         node = GraphBuilder(ctx).make_slice(kwargs, name=node.name, dtypes=out_dtypes, shapes=out_shapes)
         node = ctx.get_node_by_output(node)
         if needs_squeeze:
-            squeeze_node = ctx.insert_new_node_on_output("Squeeze", node.output[0], node.child_name())
-            squeeze_node.set_attr("axes", needs_squeeze)
+            squeeze_node = GraphBuilder(ctx).make_squeeze(
+                {"axes": needs_squeeze, "data": node.output[0]}, name=node.child_name(), return_node=True)
+            ctx.insert_node_on_output(squeeze_node, node.output[0])
             input_dtype = ctx.get_dtype(node.output[0])
             ctx.set_dtype(squeeze_node.output[0], input_dtype)
             ctx.copy_shape(node.output[0], squeeze_node.output[0])
@@ -1348,7 +1351,8 @@ def any_version(cls, opset, ctx, node, **kwargs):
             # NHWC TO CNHW, so onnx op will work on "N" which is the same as tensorflow
             if len(input_shape) == 3:
                 # insert automatically an Unsqueeze op if the input is 3d
-                unsqz1 = ctx.make_node("Unsqueeze", input_tensor.output, {"axes": [3]})
+                unsqz1 = GraphBuilder(ctx).make_unsqueeze(
+                    {"axes": [3], "data": input_tensor.output[0]}, return_node=True)
                 trans1 = ctx.make_node("Transpose", unsqz1.output, {"perm": [3, 0, 1, 2]})
             else:
                 trans1 = ctx.make_node("Transpose", input_tensor.output, {"perm": [3, 0, 1, 2]})
@@ -1377,22 +1381,21 @@ def any_version(cls, opset, ctx, node, **kwargs):
                 kwargs = {**inputs_map}
                 ctx.remove_node(node.name)
                 slice1 = GraphBuilder(ctx).make_slice(kwargs)
-                ctx.make_node("Squeeze", [slice1], {"axes": [3]},
-                              outputs=node.output, name=node.name, dtypes=dtypes, shapes=shapes)
+                GraphBuilder(ctx).make_squeeze(
+                    {"axes": [3], "data": slice1, "outputs": node.output}, name=node.name, dtypes=dtypes, shapes=shapes)
             else:
                 kwargs = {**inputs_map, "outputs": node.output}
                 ctx.remove_node(node.name)
                 GraphBuilder(ctx).make_slice(kwargs, name=node.name, dtypes=dtypes, shapes=shapes)
         else:
             def mknode(optype, inputs, attrs=None):
                 nodename = utils.make_name(node.name + '_' + optype.lower())
-                if opset < 13 or optype != 'Squeeze':
-                    return ctx.make_node(optype, inputs, attrs, name=nodename)
-                inputs.append(attrs['axes'])
-                attrs = attrs.copy()
-                attrs.pop('axes')
+                if opset >= 13 and optype == 'Squeeze':
+                    return GraphBuilder(ctx).make_squeeze(
+                        {"axes": attrs['axes'], "data": inputs[0]}, name=nodename, return_node=True)
                 return ctx.make_node(optype, inputs, attrs, name=nodename)
 
+
             # support non 3D/4D tensors and dynamic crop vals
             # dynamic slice starts at opset 10
             utils.make_sure(ctx.opset >= 11, 'non-4D tensor or non-const crops require opset 11')

tf2onnx/optimizer/back_to_back_optimizer.py

@@ -161,11 +161,20 @@ def _optimize_squeeze_unsqueeze(g, node, consumer_nodes):
         if node2.type != 'Unsqueeze':
             return []
 
-        axis1 = node.get_attr('axes').ints
-        axis2 = node2.get_attr('axes').ints
+        axes_match = False
+        if g.opset <= 12 and node.get_attr('axes').ints == node2.get_attr('axes').ints:
+            axes_match = True
+
+        # In opset 13, axes is an input. Optional for squeeze op.
+        if g.opset >= 13 and len(node.input) == 2:
+            if node.input[1] == node2.input[1]:
+                axes_match = True
+            elif node.inputs[1].is_const() and node2.inputs[1].is_const() and \
+                node.inputs[1].get_tensor_value(as_list=True) == node2.inputs[1].get_tensor_value(as_list=True):
+                axes_match = True
 
         # if squeeze followed by unsqueeze is on diff axes, skip
-        if axis1 != axis2:
+        if not axes_match:
             return []
 
         # if unsqueeze output is graph output, skip

tf2onnx/optimizer/const_fold_optimizer.py

@@ -130,10 +130,13 @@ def _fold_unsqueeze(node, graph):
         numpy expand_dims only supports to unsqueeze one dim one time, so reshape is used to simplify the logic
         """
         const_val = node.inputs[0].get_tensor_value(as_list=False)
-        axes = list(node.get_attr("axes").ints)
-        utils.make_sure(all(axis >= 0 for axis in axes), "onnx spec says it only supports positive axis")
+        if graph.opset >= 13:
+            axes = node.inputs[1].get_tensor_value(as_list=True)
+        else:
+            axes = list(node.get_attr("axes").ints)
         shape_in = const_val.shape
         dims_out = len(shape_in) + len(axes)
+        axes = [i if i >= 0 else i + dims_out for i in axes]
         # calculate the shape of output accroding to onnx Unsqueeze's spec
         # https://github.com/onnx/onnx/blob/master/docs/Operators.md#Unsqueeze
         shape_in = iter(shape_in)

tf2onnx/optimizer/transpose_optimizer.py

@@ -587,18 +587,30 @@ def _calculate_new_attr(ori_perm, ori_squeeze_axes):
         if not self._nodes_has_single_consumer_node([trans]):
             return False
 
+        axes = None
+        # in opset 13, axes is an input not attr
         if node.get_attr("axes"):
+            axes = node.get_attr("axes").ints
+        if len(node.input) > 1 and node.inputs[1].is_const():
+            axes = node.inputs[1].get_tensor_value(as_list=True)
+
+        if axes is not None:
             # switch tran and squeeze
             # 1 switch
             self._g.replace_all_inputs(node.output[0], trans.output[0])  # ops=self._g.get_nodes()
             self._g.replace_input(node, node.input[0], trans.input[0], 0)
             self._g.replace_input(trans, trans.input[0], node.output[0], 0)
             # 2 correct attr of nodes
-            squeeze_axes = sorted(list(node.get_attr("axes").ints))
+            squeeze_axes = sorted(axes)
             trans_perm = list(trans.get_attr("perm").ints)
             new_perm, new_squeeze_axes = _calculate_new_attr(ori_perm=trans_perm, ori_squeeze_axes=squeeze_axes)
             trans.set_attr("perm", new_perm)
-            node.set_attr("axes", new_squeeze_axes)
+            if self._g.opset <= 12:
+                node.set_attr("axes", new_squeeze_axes)
+            else:
+                new_axes_np = np.array(new_squeeze_axes, dtype=np.int64)
+                new_axes_const = self._g.make_const(utils.make_name(node.inputs[1].name), new_axes_np)
+                self._g.replace_inputs(node, [node.input[0], new_axes_const.output[0]])
             # 3 set shape
             squeeze_shape = self._g.get_shape(node.output[0])
             self._g.set_shape(trans.output[0], squeeze_shape)

tf2onnx/rewriter/eye_rewriter.py

@@ -6,6 +6,7 @@
 """
 
 from onnx import onnx_pb
+from tf2onnx.graph_builder import GraphBuilder
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 
 # pylint: disable=invalid-name,unused-argument,missing-docstring, unused-variable
@@ -134,8 +135,11 @@ def rewrite_eye(g, ops):
             g.remove_node(old_output.name)
 
             # onnx op "EyeLike" need a 2D tensor, so generate it
-            num_rows = g.make_node("Unsqueeze", num_rows.output, attr={"axes": [0]})
-            num_columns = g.make_node("Unsqueeze", num_columns.output, attr={"axes": [0]})
+
+            num_rows = GraphBuilder(g).make_unsqueeze(
+                {"axes": [0], "data": num_rows.output[0]}, return_node=True)
+            num_columns = GraphBuilder(g).make_unsqueeze(
+                {"axes": [0], "data": num_columns.output[0]}, return_node=True)
             matrix_shape = g.make_node("Concat", [num_rows.output[0], num_columns.output[0]], attr={"axis": 0})
             # cast nodes added for "ConstantOfShape" in ONNX only accepts int64 data.
             matrix_shape_int64 = g.make_node("Cast", matrix_shape.output, attr={"to": onnx_pb.TensorProto.INT64})