decouple rewriter

wayuanho · wayuanho · commit f74b167ca69f · 2019-07-10T06:37:50.000Z
diff --git a/test.onnx b/test.onnx
diff --git a/tf2onnx/rewriter/__init__.py b/tf2onnx/rewriter/__init__.py
@@ -7,27 +7,38 @@
 from __future__ import unicode_literals
 
 from tf2onnx.rewriter.cond_rewriter import rewrite_cond
-from tf2onnx.rewriter.random_uniform import rewrite_random_uniform, rewrite_random_uniform_fold_const
-from tf2onnx.rewriter.leakyrelu_rewriter import rewrite_leakyrelu
-from tf2onnx.rewriter.gemm_rewriter import rewrite_gemm
+from tf2onnx.rewriter.conv2d_with_pad_rewriter import rewrite_conv2d_with_pad
+from tf2onnx.rewriter.dropout_rewriter import rewrite_dropout
 from tf2onnx.rewriter.eye_rewriter import rewrite_eye
-from tf2onnx.rewriter.thresholded_relu_rewriter import rewrite_thresholded_relu
+from tf2onnx.rewriter.flatten_rewriter import rewrite_flatten
+from tf2onnx.rewriter.gemm_rewriter import rewrite_gemm
+from tf2onnx.rewriter.leakyrelu_rewriter import rewrite_leakyrelu
+from tf2onnx.rewriter.random_normal_rewriter import rewrite_random_normal
+from tf2onnx.rewriter.random_uniform import rewrite_random_uniform, rewrite_random_uniform_fold_const
 from tf2onnx.rewriter.rnn import rewrite_single_direction_lstm, rewrite_bi_direction_lstm, \
     rewrite_single_direction_gru, rewrite_bi_direction_gru, \
     rewrite_custom_rnn_cell, rewrite_generic_loop
+from tf2onnx.rewriter.thresholded_relu_rewriter import rewrite_thresholded_relu
+from tf2onnx.rewriter.transpose_rewriter import rewrite_transpose
 
 __all__ = [
     "rewrite_cond",
+    "rewrite_conv2d_with_pad",
+    "rewrite_dropout",
+    "rewrite_eye",
+    "rewrite_flatten",
+    "rewrite_gemm",
+    "rewrite_leakyrelu",
+    "rewrite_random_normal",
     "rewrite_random_uniform",
     "rewrite_random_uniform_fold_const",
-    "rewrite_leakyrelu",
     "rewrite_thresholded_relu",
-    "rewrite_eye",
+    "rewrite_transpose",
+
     "rewrite_single_direction_lstm",
     "rewrite_bi_direction_lstm",
     "rewrite_single_direction_gru",
     "rewrite_bi_direction_gru",
     "rewrite_custom_rnn_cell",
-    "rewrite_gemm",
-    "rewrite_generic_loop"
+    "rewrite_generic_loop",
 ]
diff --git a/tf2onnx/rewriter/conv2d_with_pad_rewriter.py b/tf2onnx/rewriter/conv2d_with_pad_rewriter.py
@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.rewriter - rewrite tensorflow subgraph to onnx condv2 op with pad
+"""
+
+import numpy as np
+
+from tf2onnx import handler, logging
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+logger = logging.getLogger(__name__)
+
+
+# pylint: disable=missing-docstring
+
+
+def rewrite_conv2d_with_pad(g, ops):
+    pattern = \
+        OpTypePattern("Conv2D", name="conv", inputs=[
+            OpTypePattern("Pad", name="pad"),
+            OpTypePattern("*")
+        ])
+    matcher = GraphMatcher(pattern)
+    match_results = list(matcher.match_ops(ops))
+    for match in match_results:
+        conv = match.get_op("conv")
+        pad = match.get_op("pad")
+        paddings = pad.inputs[1]
+
+        if not paddings.is_const():
+            continue
+        mode = pad.get_attr("mode")
+        if mode:
+            mode = mode.s.decode("utf-8").lower()
+        if mode not in [None, "constant"] or len(pad.input) >= 3:
+            continue
+        # Conv2D already has a pad
+        if conv.get_attr("padding") == "SAME":
+            continue
+
+        logger.debug("merge pad [%s] into conv [%s]", pad.name, conv.name)
+        paddings_val = np.array(paddings.get_tensor_value())
+        # can't pad on batch or channel dimensions
+        if np.any(paddings_val[0]) or np.any(paddings_val[3]):
+            continue
+
+        paddings_val = paddings_val[1:3]
+        paddings_val = paddings_val.transpose().flatten()
+        g.replace_input(conv, conv.input[0], pad.input[0])
+        # convert Conv2D
+        conv.type = "Conv"
+        func, _ = handler.tf_op.find_effective_op("Conv2D")
+        func(g, conv)
+        conv.skip_conversion = True
+        conv.set_attr("auto_pad", "NOTSET")
+        conv.set_attr("pads", paddings_val)
+    return ops
diff --git a/tf2onnx/rewriter/dropout_rewriter.py b/tf2onnx/rewriter/dropout_rewriter.py
@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.rewriter - rewrite tensorflow subgraph to onnx dropout op
+"""
+
+from tf2onnx import utils
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+
+# pylint: disable=missing-docstring
+
+
+def rewrite_dropout(g, ops):
+    patterns = [
+        OpTypePattern('Mul', name='outputs', inputs=[
+            OpTypePattern('RealDiv', name="input2"),
+            OpTypePattern('Floor', inputs=[
+                OpTypePattern('Add', inputs=[
+                    OpTypePattern(None, name="input3"),
+                    OpTypePattern('RandomUniform|RandomUniformLike'),
+                ])
+            ]),
+        ]),
+        OpTypePattern("Mul", name="outputs", inputs=[
+            OpTypePattern("Mul", name="input2"),
+            OpTypePattern("Cast", inputs=[
+                OpTypePattern("GreaterEqual", inputs=[
+                    OpTypePattern("RandomUniform|RandomUniformLike"),
+                    OpTypePattern(None, name="input3")
+                ])
+            ])
+        ])
+    ]
+    for pattern in patterns:
+        matcher = GraphMatcher(pattern)
+        match_results = list(matcher.match_ops(ops))
+        for match in match_results:
+            inputs2 = match.get_op('input2')
+            outputs = match.get_op('outputs')
+            op_name = utils.make_name("Dropout")
+            out_name = utils.port_name(op_name)
+            new_node = g.make_node(
+                "Dropout",
+                [inputs2.input[0]],
+                outputs=[out_name],
+                name=op_name,
+                attr={"ratio": 1.0},
+                shapes=[g.get_shape(inputs2.input[0])],
+                dtypes=[g.get_dtype(inputs2.input[0])]
+            )
+            g.replace_all_inputs(ops, outputs.output[0], new_node.output[0])
+            g.safe_remove_nodes(match.get_nodes())
+
+    # remove dropout if its ratio is 1.0
+    for node in g.get_nodes():
+        if node.type == "Dropout" and node.get_attr("ratio").f == 1.0:
+            g.replace_all_inputs(g.get_nodes(), node.output[0], node.input[0])
+            g.remove_node(node.name)
+
+    return ops
diff --git a/tf2onnx/rewriter/flatten_rewriter.py b/tf2onnx/rewriter/flatten_rewriter.py
@@ -0,0 +1,94 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.rewriter - rewrite tensorflow subgraph to onnx flatten op
+"""
+
+import numpy as np
+
+from tf2onnx import utils
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+
+# pylint: disable=missing-docstring
+
+
+def rewrite_flatten(g, ops):
+    pattern_fixed_shape_input = \
+        OpTypePattern('Reshape', name='reshape', inputs=[
+            OpTypePattern("*", name="input"),
+            OpTypePattern('Pack', name="pack", inputs=[
+                OpTypePattern('StridedSlice', name="slice", inputs=[
+                    "*", "*", "*", "*",
+                ]),
+                "*",
+            ]),
+        ])
+    pattern_non_fixed_shape_input = \
+        OpTypePattern('Reshape', name='reshape', inputs=[
+            OpTypePattern("*", name="input"),
+            OpTypePattern('Pack', name="pack", inputs=[
+                OpTypePattern('StridedSlice', name="slice", inputs=[
+                    OpTypePattern('Shape', inputs=[
+                        OpTypePattern("*", name="input2")
+                    ]),
+                    "*", "*", "*",
+                ]),
+                "*",
+            ]),
+        ])
+    matcher = GraphMatcher(pattern_fixed_shape_input)
+    match_results_1 = list(matcher.match_ops(ops))
+
+    matcher = GraphMatcher(pattern_non_fixed_shape_input)
+    match_results_2 = list(matcher.match_ops(ops))
+
+    match_results = [(match_results_1, True), (match_results_2, False)]
+    for match_results, check_fixed_input_shape in match_results:
+        for match in match_results:
+            input_node = match.get_op('input')
+            reshape_node = match.get_op('reshape')
+            pack_node = match.get_op('pack')
+            slice_node = match.get_op('slice')
+            need_rewrite = pack_node.inputs[1].is_const() and pack_node.inputs[1].get_tensor_value() == -1
+            if not need_rewrite:
+                continue
+
+            input_shape = g.get_shape(reshape_node.input[0])
+            need_rewrite = input_shape is not None
+            if not need_rewrite:
+                continue
+
+            if check_fixed_input_shape:
+                need_rewrite = slice_node.inputs[0].is_const() and \
+                               np.array_equal(list(input_shape), list(slice_node.inputs[0].get_tensor_value()))
+                if not need_rewrite:
+                    continue
+
+            begin = slice_node.inputs[1].get_tensor_value(as_list=False)
+            end = slice_node.inputs[2].get_tensor_value(as_list=False)
+            strides = slice_node.inputs[3].get_tensor_value(as_list=False)
+            need_rewrite = np.array_equal(begin, [0]) and len(end) == 1 and \
+                           np.array_equal(strides, [1]) and end[0] - begin[0] == 1
+            if not need_rewrite:
+                continue
+
+            op_name = utils.make_name("Flatten")
+            out_name = utils.port_name(op_name)
+            g.make_node("Flatten", [reshape_node.input[0]], outputs=[out_name], name=op_name)
+
+            last_dim = input_shape[-1]
+            sec_last_dim = input_shape[-2]
+            new_dim = None
+            if last_dim > 0 and sec_last_dim > 0:
+                new_dim = last_dim * sec_last_dim
+            else:
+                new_dim = -1
+
+            g.set_shape(out_name, input_shape[:-2] + [new_dim])
+            g.replace_all_inputs(ops, reshape_node.output[0], out_name)
+            to_delete = [n for n in match.get_nodes() if n != input_node]
+            g.safe_remove_nodes(to_delete)
+
+    return ops
diff --git a/tf2onnx/rewriter/random_normal_rewriter.py b/tf2onnx/rewriter/random_normal_rewriter.py
@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.rewriter - rewrite tensorflow subgraph to onnx random normal op
+"""
+
+from tf2onnx import utils
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+
+# pylint: disable=missing-docstring
+
+
+def rewrite_random_normal(g, ops):
+    pattern = \
+        OpTypePattern('Add', name='output', inputs=[
+            OpTypePattern('Mul', name='input2', inputs=[
+                OpTypePattern('RandomStandardNormal', name='input1', inputs=["*"]), "*"
+            ]), "*"
+        ])
+
+    matcher = GraphMatcher(pattern)
+    match_results = list(matcher.match_ops(ops))
+    for match in match_results:
+        output = match.get_op('output')
+        mean = output.inputs[1].get_tensor_value()
+        dtype = g.get_dtype(output.output[0])
+        op_name = utils.make_name("RandomNormal")
+        out_name = utils.port_name(op_name)
+
+        rn_op = match.get_op('input1')
+        if rn_op.inputs[0].type == "Shape":
+            shape_node = rn_op.inputs[0]
+            new_node = g.make_node("RandomNormalLike", [shape_node.input[0]], outputs=[out_name], name=op_name,
+                                   attr={"mean": mean, "scale": 1.0, "dtype": dtype})
+        else:
+            shape = g.get_shape(output.output[0])
+            new_node = g.make_node("RandomNormal", [], outputs=[out_name], name=op_name,
+                                   attr={"shape": shape, "mean": mean, "scale": 1.0, "dtype": dtype})
+
+        g.replace_all_inputs(ops, output.output[0], new_node.output[0])
+        g.safe_remove_nodes(match.get_nodes())
+    return ops
diff --git a/tf2onnx/rewriter/random_uniform.py b/tf2onnx/rewriter/random_uniform.py
@@ -33,8 +33,7 @@ def rewrite_random_uniform(g, ops):
         to_delete = list(set(match.get_nodes()))
         new_node = create_onnx_random_uniform_op(g, tmax, tmin, ru_op, output, to_delete)
         g.replace_all_inputs(ops, output.output[0], new_node.output[0])
-        for n in to_delete:
-            g.remove_node(n.name)
+        g.safe_remove_nodes(to_delete)
 
     return ops
 
diff --git a/tf2onnx/rewriter/transpose_rewriter.py b/tf2onnx/rewriter/transpose_rewriter.py
@@ -0,0 +1,34 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.rewriter - rewrite tensorflow transpose op
+"""
+
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+
+# pylint: disable=missing-docstring
+
+
+def rewrite_transpose(g, ops):
+    pattern = \
+        OpTypePattern('Transpose', name='output', inputs=[
+            OpTypePattern(None),
+            OpTypePattern('Sub', inputs=[
+                OpTypePattern('Sub', inputs=["*", "*"]),
+                OpTypePattern('Range', inputs=["*", "*", "*"]),
+            ]),
+        ])
+
+    matcher = GraphMatcher(pattern)
+    match_results = list(matcher.match_ops(ops))
+    for match in match_results:
+        output = match.get_op('output')
+        shape = g.get_shape(output.input[0])
+        dims = [i for i in range(len(shape) - 1, -1, -1)]
+        output.set_attr("perm", dims)
+        g.remove_input(output, output.input[1])
+        to_delete = [n for n in match.get_nodes() if n != output]
+        g.safe_remove_nodes(to_delete)
+    return ops
diff --git a/tf2onnx/tfonnx.py b/tf2onnx/tfonnx.py
