Merge pull request #492 from mindest/thresholdedrelu_opset10

implement ThresholdedRelu for opset10

Author: nbcsm

tests/common.py

@@ -33,7 +33,8 @@
     "validate_const_node",
     "group_nodes_by_type",
     "test_ms_domain",
-    "check_node_domain"
+    "check_node_domain",
+    "check_op_count"
 ]
 
 

tests/test_backend.py

@@ -2167,5 +2167,20 @@ def test_conv1d_5(self):
         w = np.array([3., 3., 3.], dtype=np.float32).reshape(3, 1, 1)
         self._conv1d_test(x_val, w)
 
+    @check_opset_min_version(10, "ThresholdedRelu")
+    def test_thresholded_relu(self):
+        # tf.keras.layers.ThresholdedReLU only supports `float32` for x
+        x_val = np.array([0.0, 1.0, -1.0, 2.0, -2.0, 0.5, -0.5, 1.5, -1.5], dtype=np.float32).reshape((3, 3))
+        theta_vals = [-1.0, -0.5, 0.0, 0.5, 1.0]
+        for theta_val in theta_vals:
+            x = tf.placeholder(x_val.dtype, x_val.shape, name=_TFINPUT)
+            t = tf.keras.layers.ThresholdedReLU(theta=theta_val)
+            x_ = t.call(x)
+            _ = tf.identity(x_, name=_TFOUTPUT)
+            self._run_test_case([_OUTPUT], {_INPUT: x_val},
+                                graph_validator=lambda g: check_op_count(g, "ThresholdedRelu", 1))
+            tf.reset_default_graph()
+
+
 if __name__ == '__main__':
     unittest_main()

tf2onnx/rewriter/__init__.py

@@ -9,6 +9,7 @@
 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.thresholded_relu_rewriter import rewrite_thresholded_relu
 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
@@ -18,6 +19,7 @@
     "rewrite_random_uniform",
     "rewrite_random_uniform_fold_const",
     "rewrite_leakyrelu",
+    "rewrite_thresholded_relu",
     "rewrite_single_direction_lstm",
     "rewrite_bi_direction_lstm",
     "rewrite_single_direction_gru",

tf2onnx/rewriter/thresholded_relu_rewriter.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+tf2onnx.rewriter - rewrite tensorflow subgraph to onnx ThresholdedRelu op
+"""
+
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+from tf2onnx.rewriter.leakyrelu_rewriter import _find_edge_name_between_nodes
+
+
+# pylint: disable=missing-docstring
+
+
+def rewrite_thresholded_relu(g, ops):
+    if g.opset < 10:
+        return ops
+
+    pattern = \
+        OpTypePattern('Mul', name='mul', inputs=[
+            OpTypePattern('Cast', name='cast', inputs=[
+                OpTypePattern('Greater', name='greater', inputs=[
+                    OpTypePattern('*', name='greater_input'),
+                    OpTypePattern('Const', name='theta')
+                ])
+            ]),
+            OpTypePattern('*', name='mul_input')
+        ])
+    matcher = GraphMatcher(pattern, allow_reorder=True)
+    match_results = list(matcher.match_ops(ops))
+
+    for match in match_results:
+        greater_node = match.get_op('greater')
+        greater_input_node = match.get_op('greater_input')
+        mul_node = match.get_op("mul")
+        mul_input_node = match.get_op('mul_input')
+        cast_node = match.get_op('cast')
+
+        greater_input_edge_name = _find_edge_name_between_nodes(greater_input_node, greater_node)
+        mul_input_edge_name = _find_edge_name_between_nodes(mul_input_node, mul_node)
+        if greater_input_edge_name == mul_input_edge_name:
+            theta = match.get_op('theta').get_tensor_value()
+            thresholded_relu = g.make_node("ThresholdedRelu", inputs=[mul_input_edge_name], attr={"alpha": theta},
+                                           shapes=[g.get_shape(mul_node.output[0])],
+                                           dtypes=[g.get_dtype(mul_node.output[0])])
+            ops.remove(greater_node)
+            ops.remove(cast_node)
+            ops.remove(mul_node)
+            ops.append(thresholded_relu)
+            g.replace_all_inputs(ops, mul_node.output[0], thresholded_relu.output[0])
+    return ops

tf2onnx/tfonnx.py

@@ -772,7 +772,7 @@ def compat_handler(ctx, node, **kwargs):
     rewriters = [rewrite_transpose, rewrite_flatten,
                  rewrite_random_uniform, rewrite_random_uniform_fold_const,
                  rewrite_random_normal, rewrite_dropout,
-                 rewrite_leakyrelu, rewrite_conv2d_with_pad,
+                 rewrite_leakyrelu, rewrite_thresholded_relu, rewrite_conv2d_with_pad,
                  rewrite_single_direction_lstm, rewrite_bi_direction_lstm,
                  rewrite_single_direction_gru, rewrite_bi_direction_gru,
                  rewrite_custom_rnn_cell, rewrite_generic_loop, rewrite_cond