implement ThresholdedRelu for opset10

Author: mindest

tests/test_backend.py

@@ -2138,5 +2138,19 @@ def test_non_max_suppression(self):
         self._run_test_case([_OUTPUT, _OUTPUT1], {_INPUT: boxes_val, _INPUT1: scores_val})
 
 
+    @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})
+            tf.reset_default_graph()
+
+
 if __name__ == '__main__':
     unittest_main()

tf2onnx/onnx_opset/math.py

@@ -399,3 +399,10 @@ def version_7(cls, ctx, node, **kwargs):
         ctx.remove_node(node.name)
         ctx.make_node(op_type="Sub", inputs=[node.input[0], mul.output[0]],
                       name=node.name, outputs=node.output, shapes=shapes, dtypes=dtypes)
+
+
+@tf_op("ThresholdedReLU", onnx_op="ThresholdedRelu")
+class ThresholdedRelu:
+    @classmethod
+    def version_10(cls, ctx, node, **kwargs):
+        pass

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.thresholdedrelu_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/thresholdedrelu_rewriter.py

@@ -0,0 +1,54 @@
+# 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()
+            # check disabled for now, tf requires theta to be non-negative, while onnx does not
+            # if theta < 0:
+            #     continue
+            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