Merge pull request #328 from zhijxu-MS/tmp_branch_for_PR

add leakyrelu rewriter and related test

Author: pengwa

tests/test_backend.py

@@ -721,11 +721,13 @@ def test_relu(self):
     @skip_caffe2_backend("fails on caffe2 with dim issue")
     @check_onnxruntime_incompatibility("Mul")
     def test_leaky_relu(self):
-        x_val = np.array([0.5, 1.0, -0.5, -1.0], dtype=np.float32).reshape((2, 2))
-        x = tf.placeholder(tf.float32, [2, 2], name=_TFINPUT)
-        x_ = tf.nn.leaky_relu(x)
-        _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        for alpha in [0.1, -0.1, 1.0, -1.0, 10.0, -10.0]:
+            x_val = 1000*np.random.random_sample([1000, 100]).astype(np.float32)
+            x = tf.placeholder(tf.float32, [None]*x_val.ndim, name=_TFINPUT)
+            x_ = tf.nn.leaky_relu(x, alpha)
+            _ = tf.identity(x_, name=_TFOUTPUT)
+            self._run_test_case([_OUTPUT], {_INPUT: x_val})
+            tf.reset_default_graph()
 
     @check_onnxruntime_incompatibility("Elu")
     def test_elu(self):

tf2onnx/graph.py

@@ -552,8 +552,7 @@ def _push_stack(stack, node, in_stack):
             stack.append(node)
             if node in in_stack:
                 raise ValueError('Graph has cycles.')
-            else:
-                in_stack[node] = True
+            in_stack[node] = True
 
         def _get_unvisited_child(g, node, not_visited):
             for child in g[node]:

tf2onnx/rewriter/__init__.py

@@ -10,6 +10,7 @@
     "cond_rewriter",
     "custom_rnn_rewriter",
     "gru_rewriter",
+    "leakyrelu_rewriter",
     "loop_rewriter",
     "loop_rewriter_base",
     "lstm_rewriter",

tf2onnx/rewriter/leakyrelu_rewriter.py

@@ -0,0 +1,57 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.rewriter - rewrite tensorflow subgraph to onnx leakyrelu op
+"""
+
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+
+# pylint: disable=missing-docstring
+
+
+def rewrite_leakyrelu(g, ops):
+    if g.opset < 6:
+        return ops
+
+    pattern = \
+        OpTypePattern('Maximum', name='max', inputs=[
+            OpTypePattern('Mul', name='mul', inputs=[
+                OpTypePattern('Const', name='alpha'),
+                OpTypePattern('*', name='mul_input'),
+            ]),
+            OpTypePattern('*', name='max_input'),
+        ])
+
+    matcher = GraphMatcher(pattern, allow_reorder=True)
+    match_results = list(matcher.match_ops(ops))
+    for match in match_results:
+        max_node = match.get_op('max')
+        max_input_node = match.get_op('max_input')
+        mul_node = match.get_op("mul")
+        mul_input_node = match.get_op('mul_input')
+
+        max_input_edge_name = _find_edge_name_between_nodes(max_input_node, max_node)
+        mul_input_edge_name = _find_edge_name_between_nodes(mul_input_node, mul_node)
+        if max_input_edge_name == mul_input_edge_name:
+            alpha = match.get_op("alpha").get_tensor_value()
+            if alpha >= 1:
+                continue
+            leakyrelu = g.make_node("LeakyRelu", inputs=[max_input_edge_name], attr={"alpha": alpha},
+                                    shapes=[g.get_shape(max_node.output[0])], dtypes=[g.get_dtype(max_node.output[0])])
+            ops.remove(max_node)
+            ops.remove(mul_node)
+            ops.append(leakyrelu)
+            g.replace_all_inputs(ops, max_node.output[0], leakyrelu.output[0])
+
+    return ops
+
+
+def _find_edge_name_between_nodes(src_node, consumer_node):
+    # find the first edge connection between two nodes.
+    for consumer_end in consumer_node.input:
+        for src_end in src_node.output:
+            if consumer_end == src_end:
+                return consumer_end
+    return None

tf2onnx/rewriter/random_uniform.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 """
-tf2onnx.rewrite - rewrite tensorflow subgraph to onnx random_uniform op
+tf2onnx.rewriter - rewrite tensorflow subgraph to onnx random_uniform op
 """
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx import utils

tf2onnx/shape_inference.py

@@ -98,8 +98,7 @@ def infer_shape_for_node(g, node):
             val = list(shape_attr.floats)
             if val:
                 raise ValueError("placeholder shape has floats value, and not scalar value")
-            else:
-                new_shape = ()
+            new_shape = ()
 
         if new_shape is not None:
             g.set_shape(node.output[0], new_shape)

tf2onnx/tfonnx.py

@@ -26,6 +26,7 @@
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 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.rnn import rewrite_bi_direction_gru
 from tf2onnx.rewriter.rnn import rewrite_custom_rnn_cell
 from tf2onnx.rewriter.rnn import rewrite_generic_loop
@@ -2475,7 +2476,7 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     # bi-directional re-writer should be placed after single directional re-writer
     rewriters = [rewrite_transpose, rewrite_flatten,
                  rewrite_random_uniform, rewrite_random_uniform_fold_const,
-                 rewrite_random_normal, rewrite_dropout,
+                 rewrite_random_normal, rewrite_dropout, rewrite_leakyrelu,
                  rewrite_single_direction_lstm, rewrite_bi_direction_lstm,
                  rewrite_single_direction_gru, rewrite_single_direction_grublock,
                  rewrite_bi_direction_gru, rewrite_logical_compare_with_equal,