[activations] fused gelu kernel

tensorflow_addons/activations/BUILD

@@ -6,12 +6,14 @@ py_library(
     name = "activations",
     srcs = [
         "__init__.py",
+        "gelu.py",
         "sparsemax.py",
     ],
-    srcs_version = "PY2AND3",
-    deps = [
+    data = [
+        "//tensorflow_addons/custom_ops/activations:_activation_ops.so",
         "//tensorflow_addons/utils",
     ],
+    srcs_version = "PY2AND3",
 )
 
 py_test(
@@ -26,3 +28,16 @@ py_test(
         ":activations",
     ],
 )
+
+py_test(
+    name = "gelu_test",
+    size = "large",
+    srcs = [
+        "gelu_test.py",
+    ],
+    main = "gelu_test.py",
+    srcs_version = "PY2AND3",
+    deps = [
+        ":activations",
+    ],
+)

tensorflow_addons/activations/README.md

@@ -1,14 +1,16 @@
 # Addons - Activations
 
 ## Maintainers
-| Submodule  | Maintainers  | Contact Info   |
-|:---------- |:------------- |:--------------|
-| sparsemax | @AndreasMadsen | [email protected] |
+| Submodule | Maintainers               | Contact Info                             |
+|:----------|:--------------------------|:-----------------------------------------|
+| gelu      | @AakashKumarNain @WindQAQ | [email protected] [email protected] |
+| sparsemax | @AndreasMadsen            | [email protected]                |
 
 ## Contents
-| Submodule | Activation  | Reference                               |
-|:----------------------- |:-------------------|:---------------|
-| sparsemax | Sparsemax | https://arxiv.org/abs/1602.02068           |
+| Submodule | Activation | Reference                        |
+|:----------|:-----------|:---------------------------------|
+| gelu      | gelu       | https://arxiv.org/abs/1606.08415 |
+| sparsemax | Sparsemax  | https://arxiv.org/abs/1602.02068 |
 
 
 ## Contribution Guidelines

tensorflow_addons/activations/__init__.py

@@ -18,4 +18,5 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow_addons.activations.gelu import gelu
 from tensorflow_addons.activations.sparsemax import sparsemax

tensorflow_addons/activations/gelu.py

@@ -0,0 +1,55 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+from tensorflow_addons.utils import keras_utils
+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"))
+
+
+@keras_utils.register_keras_custom_object
+@tf.function
+def gelu(x, approximate=True):
+    """Gaussian Error Linear Unit.
+
+    Computes gaussian error linear:
+    `0.5 * x * (1 + tanh(sqrt(2 / pi) * (x + 0.044715 * x^3)))` or
+    `x * P(X <= x) = 0.5 * x * (1 + erf(x / sqrt(2)))`, where P(X) ~ N(0, 1),
+    depending on whether approximation is enabled.
+
+    See [Gaussian Error Linear Units (GELUs)](https://arxiv.org/abs/1606.08415)
+    and [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805).
+
+    Args:
+        x: A `Tensor`. Must be one of the following types:
+            `float16`, `float32`, `float64`.
+        approximate: bool, whether to enable approximation.
+    Returns:
+        A `Tensor`. Has the same type as `x`.
+    """
+    x = tf.convert_to_tensor(x)
+    return _activation_ops_so.gelu(x, approximate)
+
+
+@tf.RegisterGradient("Gelu")
+def _gelu_grad(op, grad):
+    return _activation_ops_so.gelu_grad(grad, op.inputs[0],
+                                        op.get_attr("approximate"))

tensorflow_addons/activations/gelu_test.py

@@ -0,0 +1,106 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+
+import math
+
+import numpy as np
+import tensorflow as tf
+from tensorflow_addons.activations import gelu
+from tensorflow_addons.utils import test_utils
+
+
+def _ref_gelu(x, approximate=True):
+    x = tf.convert_to_tensor(x)
+    if approximate:
+        pi = tf.cast(math.pi, x.dtype)
+        coeff = tf.cast(0.044715, x.dtype)
+        return 0.5 * x * (
+            1.0 + tf.tanh(tf.sqrt(2.0 / pi) * (x + coeff * tf.pow(x, 3))))
+    else:
+        return 0.5 * x * (
+            1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0), x.dtype)))
+
+
+@test_utils.run_all_in_graph_and_eager_modes
+class GeluTest(tf.test.TestCase, parameterized.TestCase):
+    @parameterized.named_parameters(("float16", np.float16),
+                                    ("float32", np.float32),
+                                    ("float64", np.float64))
+    def test_gelu(self, dtype):
+        x = np.random.rand(2, 3, 4).astype(dtype)
+        self.assertAllCloseAccordingToType(gelu(x), _ref_gelu(x))
+        self.assertAllCloseAccordingToType(gelu(x, False), _ref_gelu(x, False))
+
+    @parameterized.named_parameters(("float16", np.float16),
+                                    ("float32", np.float32),
+                                    ("float64", np.float64))
+    def test_gradients(self, dtype):
+        x = tf.constant([1.0, 2.0, 3.0], dtype=dtype)
+
+        for approximate in [True, False]:
+            with self.subTest(approximate=approximate):
+                with tf.GradientTape(persistent=True) as tape:
+                    tape.watch(x)
+                    y_ref = _ref_gelu(x, approximate)
+                    y = gelu(x, approximate)
+                grad_ref = tape.gradient(y_ref, x)
+                grad = tape.gradient(y, x)
+                self.assertAllCloseAccordingToType(grad, grad_ref)
+
+    @parameterized.named_parameters(("float32", np.float32),
+                                    ("float64", np.float64))
+    def test_theoretical_gradients(self, dtype):
+        # Only test theoretical gradients for float32 and float64
+        # because of the instability of float16 while computing jacobian
+        x = tf.constant([1.0, 2.0, 3.0], dtype=dtype)
+
+        for approximate in [True, False]:
+            with self.subTest(approximate=approximate):
+                theoretical, numerical = tf.test.compute_gradient(
+                    lambda x: gelu(x, approximate=approximate), [x])
+                self.assertAllCloseAccordingToType(
+                    theoretical, numerical, atol=1e-4)
+
+    def test_unknown_shape(self):
+        fn = gelu.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), gelu(x))
+
+    def test_serialization(self):
+        ref_fn = gelu
+        config = tf.keras.activations.serialize(ref_fn)
+        fn = tf.keras.activations.deserialize(config)
+        self.assertEqual(fn, ref_fn)
+
+    def test_serialization_with_layers(self):
+        layer = tf.keras.layers.Dense(3, activation=gelu)
+        config = tf.keras.layers.serialize(layer)
+        deserialized_layer = tf.keras.layers.deserialize(config)
+        self.assertEqual(deserialized_layer.__class__.__name__,
+                         layer.__class__.__name__)
+        self.assertEqual(deserialized_layer.activation.__name__, "gelu")
+
+
+if __name__ == "__main__":
+    tf.test.main()

tensorflow_addons/custom_ops/activations/BUILD

@@ -0,0 +1,47 @@
+licenses(["notice"])  # Apache 2.0
+
+package(default_visibility = ["//visibility:public"])
+
+load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI")
+load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured", "if_cuda")
+
+cc_library(
+    name = "gelu_op_gpu",
+    srcs = [
+        "cc/kernels/gelu_op.h",
+        "cc/kernels/gelu_op_gpu.cu.cc",
+    ],
+    copts = if_cuda_is_configured([
+        "-DGOOGLE_CUDA=1",
+        "-x cuda",
+        "-nvcc_options=relaxed-constexpr",
+        "-nvcc_options=ftz=true",
+    ]),
+    deps = [
+        "@local_config_tf//:libtensorflow_framework",
+        "@local_config_tf//:tf_header_lib",
+    ] + if_cuda_is_configured([
+        "@local_config_cuda//cuda:cuda_libs",
+        "@local_config_cuda//cuda:cuda_headers",
+    ]),
+    alwayslink = 1,
+)
+
+cc_binary(
+    name = "_activation_ops.so",
+    srcs = [
+        "cc/kernels/gelu_op.cc",
+        "cc/kernels/gelu_op.h",
+        "cc/ops/gelu_op.cc",
+    ],
+    copts = [
+        "-pthread",
+        "-std=c++11",
+        D_GLIBCXX_USE_CXX11_ABI,
+    ] + if_cuda(["-DGOOGLE_CUDA=1"]),
+    linkshared = 1,
+    deps = [
+        "@local_config_tf//:libtensorflow_framework",
+        "@local_config_tf//:tf_header_lib",
+    ] + if_cuda_is_configured([":gelu_op_gpu"]),
+)

tensorflow_addons/custom_ops/activations/cc/kernels/gelu_op.cc

@@ -0,0 +1,77 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#define EIGEN_USE_THREADS
+
+#include "tensorflow_addons/custom_ops/activations/cc/kernels/gelu_op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+
+namespace tensorflow {
+
+using CPUDevice = Eigen::ThreadPoolDevice;
+
+#define REGISTER_GELU_KERNELS(type)                                  \
+  REGISTER_KERNEL_BUILDER(                                           \
+      Name("Gelu").Device(DEVICE_CPU).TypeConstraint<type>("T"),     \
+      GeluOp<CPUDevice, type>);                                      \
+  REGISTER_KERNEL_BUILDER(                                           \
+      Name("GeluGrad").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
+      GeluGradOp<CPUDevice, type>);
+
+// Gelu only makes sense with floating points.
+TF_CALL_GPU_NUMBER_TYPES(REGISTER_GELU_KERNELS);
+#undef REGISTER_GELU_KERNELS
+
+#ifdef GOOGLE_CUDA
+
+using GPUDevice = Eigen::GpuDevice;
+
+// Forward declarations of the functor specializations for GPU.
+namespace functor {
+#define DECLARE_GPU_SPEC(T)                                          \
+  template <>                                                        \
+  void Gelu<GPUDevice, T>::operator()(                               \
+      const GPUDevice& d, typename TTypes<T>::ConstTensor features,  \
+      bool approximate, typename TTypes<T>::Tensor activations);     \
+  extern template struct Gelu<GPUDevice, T>;                         \
+                                                                     \
+  template <>                                                        \
+  void GeluGrad<GPUDevice, T>::operator()(                           \
+      const GPUDevice& d, typename TTypes<T>::ConstTensor gradients, \
+      typename TTypes<T>::ConstTensor features, bool approximate,    \
+      typename TTypes<T>::Tensor backprops);                         \
+  extern template struct GeluGrad<GPUDevice, T>;
+
+TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPEC);
+#undef DECLARE_GPU_SPEC
+}  // namespace functor
+
+// Registration of the GPU implementations.
+#define REGISTER_GELU_GPU_KERNELS(type)                              \
+  REGISTER_KERNEL_BUILDER(                                           \
+      Name("Gelu").Device(DEVICE_GPU).TypeConstraint<type>("T"),     \
+      GeluOp<GPUDevice, type>);                                      \
+  REGISTER_KERNEL_BUILDER(                                           \
+      Name("GeluGrad").Device(DEVICE_GPU).TypeConstraint<type>("T"), \
+      GeluGradOp<GPUDevice, type>);
+
+TF_CALL_GPU_NUMBER_TYPES(REGISTER_GELU_GPU_KERNELS);
+#undef REGISTER_GELU_GPU_KERNELS
+
+#endif  // GOOGLE_CUDA
+
+}  // namespace tensorflow