Merge remote-tracking branch 'upstream/master' into jignparm/optimize_transpose_multiply

Author: jignparm

tests/test_backend.py

@@ -1787,6 +1787,29 @@ def func(x, y):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
 
+    @check_opset_min_version(10, "Slice")
+    def test_strided_slice_reverse_1(self):
+        x_val = np.arange(16 * 32).astype(np.float32).reshape((1, 16, 32, 1))
+        def func(x):
+            return tf.concat([x[:, :, :10], x[:, :, :21:-1]], axis=0, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
+    @check_opset_min_version(10, "Slice")
+    def test_strided_slice_reverse_2(self):
+        x_val = np.arange(16 * 32).astype(np.float32).reshape((1, 16, 32, 1))
+        def func(x):
+            return tf.concat([x[:, :, :10], x[:, :, 9::-1]], axis=0, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
+    @check_opset_min_version(10, "Slice")
+    def test_strided_slice_reverse_3(self):
+        x_val = np.zeros((1, 16, 32, 1)).astype(np.float32)
+        y_val = np.array(9).astype(np.int32)
+        z_val = np.array(-1).astype(np.int32)
+        def func(x, y, z):
+            return tf.concat([x[:, :, :10], x[:, :, y::z]], axis=0, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val, _INPUT2: z_val})
+
     @check_opset_min_version(10, "Slice")
     def test_new_axis_mask(self):
         def func(x, y):

tf2onnx/constants.py

@@ -37,3 +37,9 @@
 
 # Environment variables
 ENV_TF2ONNX_DEBUG_MODE = "TF2ONNX_DEBUG_MODE"
+
+# Mapping opset to IR version.
+# When adding here, make sure that the IR changes don't impact that we do.
+OPSET_TO_IR_VERSION = {
+    1: 3, 2:3, 3: 3, 4: 4, 5: 3, 6:3, 7:3, 8:4, 9:4, 10:5, 11:6, 12:7
+}

tf2onnx/custom_opsets/ms.py

@@ -92,3 +92,14 @@ def version_1(cls, ctx, node, **kwargs):
         node.attr.pop("padding")
         if "explicit_paddings" in node.attr:
             node.attr.pop("explicit_paddings")
+
+@tf_op("CropAndResize", domain=constants.MICROSOFT_DOMAIN)
+class CropAndResize:
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+        """ utilize contrib cropandresize """
+        node.attr['method'].name = 'mode'
+        node.domain = constants.MICROSOFT_DOMAIN
+        ctx.insert_new_node_on_input(node, "Transpose", node.input[0], perm=constants.NHWC_TO_NCHW)
+        ctx.insert_new_node_on_output("Transpose", node.output[0], node.name + '_transposed',
+                                      None, perm=constants.NCHW_TO_NHWC)

tf2onnx/graph.py

@@ -1039,6 +1039,12 @@ def make_model(self, graph_doc, optimize=False, graph_name="tf2onnx", **kwargs):
             kwargs["opset_imports"] = opsets
         model_proto = helper.make_model(graph, **kwargs)
 
+        # set the IR version based on opset
+        try:
+            model_proto.ir_version = constants.OPSET_TO_IR_VERSION.get(self.opset, model_proto.ir_version)
+        except:
+            logger.error("ir_version override failed - install the latest onnx version")
+
         # optimize the model proto.
         # TODO: this is disabled by default because of bugs in fuse_consecutive_transposes
         if optimize:

tf2onnx/onnx_opset/controlflow.py

@@ -372,14 +372,14 @@ def version_9(cls, ctx, node, **kwargs):
         ctx.copy_dtype(node.output[0], transpose_node.output[0])
 
 
-@tf_op("IteratorV2")
+@tf_op("IteratorV2", "FIFOQueueV2")
 class Iterator:
     @classmethod
     def version_8(cls, ctx, node, **kwargs):
         ctx.remove_node(node.name)
 
 
-@tf_op("IteratorGetNext")
+@tf_op("IteratorGetNext", "QueueDequeueV2")
 class IteratorGetNext:
     @classmethod
     def version_8(cls, ctx, node, **kwargs):

tf2onnx/onnx_opset/math.py

@@ -97,8 +97,8 @@ def make_min_or_max_op(ctx, op_type, inputs, outputs,
         ctx.set_dtype(cast_node.output[0], origin_dtype)
         ctx.copy_shape(node.output[0], cast_node.output[0])
         actual_outputs = cast_node.output
-    ctx.make_node("Identity", actual_outputs, outputs=outputs,
-                  shapes=output_shapes, dtypes=output_dtypes)
+    final_node = ctx.make_node("Identity", actual_outputs, outputs=outputs,
+                               shapes=output_shapes, dtypes=output_dtypes)
 
     # tensorflow minimum/maximum does support broadcast, onnx < opset 8 does not.
     # handle this by doing something like:
@@ -124,6 +124,7 @@ def make_min_or_max_op(ctx, op_type, inputs, outputs,
             add_node = ctx.make_node("Add", [input_node.output[0], sub_node.output[0]],
                                      op_name_scope=input_node.name)
             node.input[i] = add_node.output[0]
+    return final_node
 
 
 @tf_op("Minimum", onnx_op="Min")

tf2onnx/onnx_opset/tensor.py

@@ -845,37 +845,51 @@ def version_10(cls, ctx, node, **kwargs):
         # mask begin
         new_begin_mask = np.array(new_begin_mask, dtype=np_dtype)
         if not np.all(new_begin_mask == 1):
-            if begin.is_const():
-                begin = ctx.make_const(
-                    utils.make_name("begin_masked"),
-                    begin.get_tensor_value(as_list=False) * new_begin_mask
-                )
+            if begin.is_const() and strides.is_const():
+                new_begin_vals = np.copy(begin.get_tensor_value(as_list=False))
+                strides_vals = strides.get_tensor_value(as_list=False)
+                idx1 = np.where(new_begin_mask == 0)
+                idx2 = np.where(strides_vals < 0)
+                idx3 = np.intersect1d(idx1, idx2)
+                new_begin_vals[idx3] = max_size
+                begin = ctx.make_const(utils.make_name("begin_masked"), new_begin_vals)
             else:
-                begin_mask_const = ctx.make_const(
-                    utils.make_name("begin_mask"),
-                    new_begin_mask
-                )
-                begin = ctx.make_node(
-                    "Mul", [begin.output[0], begin_mask_const.output[0]],
-                    op_name_scope=node.name
-                )
+                begin_mask_const = ctx.make_const(utils.make_name("begin_mask"), np.equal(new_begin_mask, 0))
+                zero_const = ctx.make_const(utils.make_name("zero_const"), np.zeros(1, dtype=np_dtype))
+                max_const = ctx.make_const(utils.make_name("max_const"), np.array(max_size, dtype=np_dtype))
+                op1 = ctx.make_node("Less", [strides.output[0], zero_const.output[0]], op_name_scope=node.name)
+                op2 = ctx.make_node("And", [op1.output[0], begin_mask_const.output[0]], op_name_scope=node.name)
+                begin = ctx.make_node("Where", [op2.output[0], max_const.output[0], begin.output[0]],
+                                      op_name_scope=node.name)
+
         # mask end
         new_end_mask = np.array(new_end_mask, dtype=np_dtype)
         end_output = end.output[0]
         if not np.all(new_end_mask == min_size):
-            if end.is_const():
-                end = ctx.make_const(
-                    utils.make_name("end_masked"),
-                    np.maximum(end.get_tensor_value(as_list=False), new_end_mask)
-                )
+            if end.is_const() and strides.is_const():
+                new_end_mask = np.maximum(end.get_tensor_value(as_list=False), new_end_mask)
+                idx = np.where(new_end_mask == max_size)
+                sign = np.sign(strides.get_tensor_value(as_list=False))[idx]
+                new_end_mask[idx] = new_end_mask[idx] * sign
+                end = ctx.make_const(utils.make_name("end_masked"), new_end_mask)
                 end_output = end.output[0]
             else:
-                end_mask_const = ctx.make_const(
-                    utils.make_name("end_mask"),
-                    np.array(new_end_mask, dtype=np_dtype)
-                )
-                end_output = utils.make_name("{}__end".format(node.name))
-                math.make_min_or_max_op(ctx, "Max", [end.output[0], end_mask_const.output[0]], [end_output])
+                # Overlay new_end_mask with specified end values.
+                # Adjust max_size to min_size if steps are < 0
+                max_const = ctx.make_const(utils.make_name("max_const"), np.array(max_size, dtype=np_dtype))
+                min_const = ctx.make_const(utils.make_name("min_const"), np.array(min_size, dtype=np_dtype))
+                zero_const = ctx.make_const(utils.make_name("zero_const"), np.zeros(1, dtype=np_dtype))
+                end_mask_const = ctx.make_const(utils.make_name("end_mask"), np.array(new_end_mask, dtype=np_dtype))
+                outputname = utils.make_name("{}__newendmask".format(node.name))
+                new_end_mask = math.make_min_or_max_op(ctx, "Max", [end.output[0], end_mask_const.output[0]],
+                                                       [outputname])
+                op1 = ctx.make_node("Less", [strides.output[0], zero_const.output[0]], op_name_scope=node.name)
+                op2 = ctx.make_node("Equal", [new_end_mask.output[0], max_const.output[0]], op_name_scope=node.name)
+                op3 = ctx.make_node("And", [op2.output[0], op1.output[0]], op_name_scope=node.name)
+                final_end = ctx.make_node("Where", [op3.output[0], min_const.output[0],
+                                                    new_end_mask.output[0]], op_name_scope=node.name)
+                end_output = final_end.output[0]
+
         # mask strides for shrink
         shrink_strided_mask = np.array(shrink_strided_mask, dtype=np_dtype)
         strides_output = strides.output[0]

tf2onnx/rewriter/gemm_rewriter.py

@@ -99,4 +99,16 @@ def get_gemm_attr(match):
                     return attr, False
                 match_args = match_args[0]
             attr[arg] = match_args
+    for arg in ["matmul"]:
+        arg_op = match.get_op(arg)
+        if arg_op is not None:
+            match_args = arg_op.attr
+            if isinstance(match_args, dict):
+                keys = list(match_args.keys())
+                if 'transpose_a' not in keys and 'transpose_b' not in keys:
+                    return attr, False
+                match_args_a = match_args['transpose_a'].i
+                attr['transA'] = match_args_a
+                match_args_b = match_args['transpose_b'].i
+                attr['transB'] = match_args_b
     return attr, True

tf2onnx/tf_utils.py

@@ -138,7 +138,7 @@ def tflist_to_onnx(g, shape_override):
                     "Tout", "Tlabels", "Tindex", "element_shape", "Targmax", "Tperm", "Tcond",
                     "T_threshold", "element_dtype", "shape_type", "_lower_using_switch_merge",
                     "parallel_iterations", "_num_original_outputs", "output_types", "output_shapes",
-                    "key_dtype", "value_dtype", "Tin", "Tout"]
+                    "key_dtype", "value_dtype", "Tin", "Tout", "capacity", "component_types", "shapes"]
 
     node_list = g.get_operations()
     functions = {}