Thresholded Linear Unit

tensorflow_addons/layers/BUILD

@@ -12,6 +12,7 @@ py_library(
         "optical_flow.py",
         "poincare.py",
         "sparsemax.py",
+        "tlu.py",
         "wrappers.py",
     ],
     data = [
@@ -35,6 +36,18 @@ py_test(
     ],
 )
 
+py_test(
+    name = "tlu_test",
+    size = "small",
+    srcs = [
+        "tlu_test.py",
+    ],
+    main = "tlu_test.py",
+    deps = [
+        ":layers",
+    ],
+)
+
 py_test(
     name = "layers_wrappers_test",
     size = "small",

tensorflow_addons/layers/__init__.py

@@ -25,4 +25,5 @@
 from tensorflow_addons.layers.optical_flow import CorrelationCost
 from tensorflow_addons.layers.poincare import PoincareNormalize
 from tensorflow_addons.layers.sparsemax import Sparsemax
+from tensorflow_addons.layers.tlu import TLU
 from tensorflow_addons.layers.wrappers import WeightNormalization

tensorflow_addons/layers/tlu.py

@@ -0,0 +1,120 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Implements Thresholded Linear Unit."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+@tf.keras.utils.register_keras_serializable(package='Addons')
+class TLU(tf.keras.layers.Layer):
+    """Thresholded Linear Unit. An activation function which is similar to ReLU
+    but with a learned threshold that benefits models using FRN(Filter Response
+    Normalization). Original paper: https://arxiv.org/pdf/1911.09737.
+
+    Input shape:
+        Arbitrary. Use the keyword argument `input_shape`
+        (tuple of integers, does not include the samples axis)
+        when using this layer as the first layer in a model.
+
+    Output shape:
+        Same shape as the input.
+
+    Arguments:
+        affine: bool. Whether to make it TLU-Affine or not
+        which has the form `max(x, alpha*x + tau)`
+    """
+
+    def __init__(self,
+                 affine=False,
+                 tau_initializer='zeros',
+                 tau_regularizer=None,
+                 tau_constraint=None,
+                 alpha_initializer='zeros',
+                 alpha_regularizer=None,
+                 alpha_constraint=None,
+                 **kwargs):
+        super(TLU, self).__init__(**kwargs)
+        self.supports_masking = True
+        self.affine = affine
+        self.tau_initializer = tf.keras.initializers.get(tau_initializer)
+        self.tau_regularizer = tf.keras.regularizers.get(tau_regularizer)
+        self.tau_constraint = tf.keras.constraints.get(tau_constraint)
+        if self.affine:
+            self.alpha_initializer = tf.keras.initializers.get(
+                alpha_initializer)
+            self.alpha_regularizer = tf.keras.regularizers.get(
+                alpha_regularizer)
+            self.alpha_constraint = tf.keras.constraints.get(alpha_constraint)
+
+    def build(self, input_shape):
+        param_shape = list(input_shape[1:])
+        self.tau = self.add_weight(
+            shape=param_shape,
+            name='tau',
+            initializer=self.tau_initializer,
+            regularizer=self.tau_regularizer,
+            constraint=self.tau_constraint,
+            synchronization=tf.VariableSynchronization.AUTO,
+            aggregation=tf.VariableAggregation.MEAN)
+        if self.affine:
+            self.alpha = self.add_weight(
+                shape=param_shape,
+                name='alpha',
+                initializer=self.alpha_initializer,
+                regularizer=self.alpha_regularizer,
+                constraint=self.alpha_constraint,
+                synchronization=tf.VariableSynchronization.AUTO,
+                aggregation=tf.VariableAggregation.MEAN)
+
+        axes = {i: input_shape[i] for i in range(1, len(input_shape))}
+        self.input_spec = tf.keras.layers.InputSpec(
+            ndim=len(input_shape), axes=axes)
+        self.built = True
+
+    def call(self, inputs):
+        if self.affine:
+            return tf.maximum(inputs, self.alpha * inputs + self.tau)
+        else:
+            return tf.maximum(inputs, self.tau)
+
+    def get_config(self):
+        config = {
+            'tau_initializer':
+            tf.keras.initializers.serialize(self.tau_initializer),
+            'tau_regularizer':
+            tf.keras.regularizers.serialize(self.tau_regularizer),
+            'tau_constraint':
+            tf.keras.constraints.serialize(self.tau_constraint),
+            'affine':
+            self.affine
+        }
+
+        if self.affine:
+            config['alpha_initializer'] = tf.keras.initializers.serialize(
+                self.alpha_initializer)
+            config['alpha_regularizer'] = tf.keras.regularizers.serialize(
+                self.alpha_regularizer)
+            config['alpha_constraint'] = tf.keras.constraints.serialize(
+                self.alpha_constraint)
+
+        base_config = super(TLU, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    def compute_output_shape(self, input_shape):
+        return input_shape

tensorflow_addons/layers/tlu_test.py

@@ -0,0 +1,63 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Tests for TLU activation."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow_addons.layers.tlu import TLU
+from tensorflow_addons.utils import test_utils
+
+
+@parameterized.parameters([np.float16, np.float32, np.float64])
+@test_utils.run_all_in_graph_and_eager_modes
+class TestTLU(tf.test.TestCase):
+    def test_random(self, dtype):
+        x = np.array([[-2.5, 0., 0.3]]).astype(dtype)
+        val = np.array([[0., 0., 0.3]]).astype(dtype)
+        test_utils.layer_test(
+            TLU, kwargs={'dtype': dtype}, input_data=x, expected_output=val)
+
+    def test_affine(self, dtype):
+        x = np.array([[-2.5, 0., 0.3]]).astype(dtype)
+        val = np.array([[-1.5, 1.0, 1.3]]).astype(dtype)
+        test_utils.layer_test(
+            TLU,
+            kwargs={
+                'affine': True,
+                'dtype': dtype,
+                'alpha_initializer': 'ones',
+                'tau_initializer': 'ones'
+            },
+            input_data=x,
+            expected_output=val)
+
+    def test_serialization(self, dtype):
+        tlu = TLU(
+            affine=True,
+            alpha_initializer='ones',
+            tau_initializer='ones',
+            dtype=dtype)
+        serialized_tlu = tf.keras.layers.serialize(tlu)
+        new_layer = tf.keras.layers.deserialize(serialized_tlu)
+        self.assertEqual(tlu.get_config(), new_layer.get_config())
+
+
+if __name__ == '__main__':
+    tf.test.main()