Rrelu fix (#820)

* refactor rrelu

Author: fsx950223

tensorflow_addons/activations/BUILD

@@ -123,7 +123,7 @@ py_test(
     srcs = [
         "rrelu_test.py",
     ],
-    flaky = True,
+    args = ["--benchmarks=all"],
     main = "rrelu_test.py",
     deps = [
         ":activations",

tensorflow_addons/activations/rrelu.py

@@ -18,14 +18,9 @@
 from __future__ import print_function
 
 import tensorflow as tf
-from tensorflow_addons.utils.resource_loader import get_path_to_datafile
-
-_activation_ops_so = tf.load_op_library(
-    get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
 
 
 @tf.keras.utils.register_keras_serializable(package='Addons')
-@tf.function
 def rrelu(x, lower=0.125, upper=0.3333333333333333, training=None, seed=None):
     """rrelu function.
 
@@ -51,14 +46,11 @@ def rrelu(x, lower=0.125, upper=0.3333333333333333, training=None, seed=None):
     if training is None:
         training = tf.keras.backend.learning_phase()
         training = bool(tf.keras.backend.get_value(training))
-    # TODO: get rid of v1 API
-    seed1, seed2 = tf.compat.v1.random.get_seed(seed)
-    result, _ = _activation_ops_so.addons_rrelu(x, lower, upper, training,
-                                                seed1, seed2)
-    return result
 
+    if training:
+        alpha = tf.random.uniform(
+            tf.shape(x), minval=lower, maxval=upper, dtype=x.dtype, seed=seed)
+    else:
+        alpha = tf.cast((lower + upper) / 2, x.dtype)
 
-@tf.RegisterGradient("Addons>Rrelu")
-def _rrelu_grad(op, *grad):
-    return _activation_ops_so.addons_rrelu_grad(grad[0], op.inputs[0],
-                                                op.outputs[1])
+    return tf.where(x >= 0, x, alpha * x)

tensorflow_addons/activations/rrelu_test.py

@@ -24,54 +24,68 @@
 from tensorflow_addons.activations import rrelu
 from tensorflow_addons.utils import test_utils
 
-
-def _ref_rrelu(x, lower, upper):
-    return tf.where(x >= 0, x, (lower + upper) * x / 2)
+SEED = 111111
 
 
 @test_utils.run_all_in_graph_and_eager_modes
 class RreluTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(("float16", np.float16),
                                     ("float32", np.float32),
                                     ("float64", np.float64))
-    @tf.function
     def test_rrelu(self, dtype):
         x = tf.constant([-2.0, -1.0, 0.0, 1.0, 2.0], dtype=dtype)
         lower = 0.1
         upper = 0.2
-        result = rrelu(x, lower, upper, training=False)
-        expect_result = _ref_rrelu(x, lower, upper)
-        self.assertAllCloseAccordingToType(result, expect_result)
+
+        training_results = {
+            np.float16: [-0.288330078, -0.124206543, 0, 1, 2],
+            np.float32: [-0.26851666, -0.116421416, 0, 1, 2],
+            np.float64: [-0.3481333923206531, -0.17150176242558851, 0, 1, 2],
+        }
+        for training in [True, False]:
+            with self.subTest(training=training):
+                tf.random.set_seed(SEED)
+                result = rrelu(x, lower, upper, training=training, seed=SEED)
+                if training:
+                    expect_result = training_results.get(dtype)
+                else:
+                    expect_result = [
+                        -0.30000001192092896, -0.15000000596046448, 0, 1, 2
+                    ]
+                self.assertAllCloseAccordingToType(result, expect_result)
 
     @parameterized.named_parameters(("float32", np.float32),
                                     ("float64", np.float64))
     def test_theoretical_gradients(self, dtype):
-        x = tf.constant([-2.0, -1.0, -0.1, 0.1, 1.0, 2.0], dtype=dtype)
-        lower = 0.1
-        upper = 0.2
-        for training in [True, False]:
-            with self.subTest(training=training):
-                theoretical, numerical = tf.test.compute_gradient(
-                    lambda x: rrelu(
-                        x, lower, upper, training=training, seed=111111), [x])
-                # TODO: investigate the difference between CPU and GPU
-                if training is True and tf.test.is_gpu_available() is False:
-                    numerical = [[[0.134971, 0., 0., 0., 0., 0.],
-                                  [0., 0.15648358, 0., 0., 0., 0.],
-                                  [0., 0., 0.18776372, 0., 0., 0.],
-                                  [0., 0., 0., 1., 0., 0.],
-                                  [0., 0., 0., 0., 1., 0.],
-                                  [0., 0., 0., 0., 0., 1.]]]
-                self.assertAllCloseAccordingToType(
-                    theoretical, numerical, rtol=5e-4, atol=5e-4)
-
-    def test_unknown_shape(self):
-        fn = rrelu.get_concrete_function(
-            tf.TensorSpec(shape=None, dtype=tf.float32))
-
-        for shape in [(1,), (1, 2), (1, 2, 3), (1, 2, 3, 4)]:
-            x = tf.ones(shape=shape, dtype=tf.float32)
-            self.assertAllClose(fn(x), rrelu(x))
+        if tf.executing_eagerly():
+
+            def rrelu_wrapper(lower, upper, training):
+                def inner(x):
+                    tf.random.set_seed(SEED)
+                    return rrelu(x, lower, upper, training=training, seed=SEED)
+
+                return inner
+
+            x = tf.constant([-2.0, -1.0, -0.1, 0.1, 1.0, 2.0], dtype=dtype)
+            lower = 0.1
+            upper = 0.2
+
+            for training in [True, False]:
+                with self.subTest(training=training):
+                    theoretical, numerical = tf.test.compute_gradient(
+                        rrelu_wrapper(lower, upper, training), [x])
+                    self.assertAllCloseAccordingToType(
+                        theoretical, numerical, rtol=5e-4, atol=5e-4)
+
+
+class RreluBenchmarks(tf.test.Benchmark):
+    def benchmarkRreluOp(self):
+        with tf.compat.v1.Session(config=tf.test.benchmark_config()) as sess:
+            x = tf.constant([-2.0, -1.0, 0.0, 1.0, 2.0], dtype=np.float32)
+            lower = 0.1
+            upper = 0.2
+            result = rrelu(x, lower, upper, training=True)
+            self.run_op_benchmark(sess, result.op, min_iters=25)
 
 
 if __name__ == "__main__":

tensorflow_addons/custom_ops/activations/BUILD

@@ -15,8 +15,6 @@ custom_op_library(
         "cc/kernels/lisht_op.h",
         "cc/kernels/mish_op.cc",
         "cc/kernels/mish_op.h",
-        "cc/kernels/rrelu_op.cc",
-        "cc/kernels/rrelu_op.h",
         "cc/kernels/softshrink_op.cc",
         "cc/kernels/softshrink_op.h",
         "cc/kernels/tanhshrink_op.cc",
@@ -25,7 +23,6 @@ custom_op_library(
         "cc/ops/hardshrink_op.cc",
         "cc/ops/lisht_op.cc",
         "cc/ops/mish_op.cc",
-        "cc/ops/rrelu_op.cc",
         "cc/ops/softshrink_op.cc",
         "cc/ops/tanhshrink_op.cc",
     ],
@@ -38,8 +35,6 @@ custom_op_library(
         "cc/kernels/lisht_op_gpu.cu.cc",
         "cc/kernels/mish_op.h",
         "cc/kernels/mish_op_gpu.cu.cc",
-        "cc/kernels/rrelu_op.h",
-        "cc/kernels/rrelu_op_gpu.cu.cc",
         "cc/kernels/softshrink_op.h",
         "cc/kernels/softshrink_op_gpu.cu.cc",
         "cc/kernels/tanhshrink_op.h",