From 23c01eaf7a364aee4446ec55a838163027f0ac62 Mon Sep 17 00:00:00 2001 From: Prannaya Date: Sat, 9 Sep 2023 15:13:52 +0800 Subject: [PATCH] feat(activations): Implement all activation functions in TFA --- kca/activations/__init__.py | 20 +++++++--- kca/activations/hardshrink.py | 43 ++++++++++++++++++++ kca/activations/lisht.py | 28 +++++++++++++ kca/activations/mish.py | 28 +++++++++++++ kca/activations/rrelu.py | 75 +++++++++++++++++++++++++++++++++++ kca/activations/snake.py | 30 ++++++++++++++ kca/activations/softshrink.py | 41 +++++++++++++++++++ kca/activations/tanhshrink.py | 26 ++++++++++++ kca/utils/types.py | 6 ++- 9 files changed, 291 insertions(+), 6 deletions(-) create mode 100644 kca/activations/hardshrink.py create mode 100644 kca/activations/lisht.py create mode 100644 kca/activations/mish.py create mode 100644 kca/activations/rrelu.py create mode 100644 kca/activations/snake.py create mode 100644 kca/activations/softshrink.py create mode 100644 kca/activations/tanhshrink.py diff --git a/kca/activations/__init__.py b/kca/activations/__init__.py index d643ba5..6d859b8 100644 --- a/kca/activations/__init__.py +++ b/kca/activations/__init__.py @@ -1,10 +1,20 @@ # activations/__init__.py __all__ = [ - "glu", "reglu", "geglu", "swiglu", "seglu", - "sparsemax", "differentiable_binary" + "glu", "reglu", "geglu", "swiglu", "seglu", "rrelu", + "hardshrink", "softshrink", "tanhshrink", + "lisht", "mish", "snake", "sparsemax", "differentiable_binary" ] -from kca.activations.glu import glu, reglu, geglu, swiglu, seglu -from kca.activations.sparsemax import sparsemax -from kca.activations.differentiable_binary import differentiable_binary +from .glu import glu, reglu, geglu, swiglu, seglu +from .rrelu import rrelu + +from .hardshrink import hardshrink +from .softshrink import softshrink +from .tanhshrink import tanhshrink + +from .lisht import lisht +from .mish import mish +from .snake import snake +from .sparsemax import sparsemax +from .differentiable_binary import differentiable_binary diff --git a/kca/activations/hardshrink.py b/kca/activations/hardshrink.py new file mode 100644 index 0000000..1f20d25 --- /dev/null +++ b/kca/activations/hardshrink.py @@ -0,0 +1,43 @@ +from keras_core import ops + +from kca.utils.types import TensorLike, Number + + +def hardshrink(x: TensorLike, lower: Number = -0.5, upper: Number = 0.5) -> TensorLike: + r"""Hard shrink function. + + Computes hard shrink function: + + $$ + \mathrm{hardshrink}(x) = + \begin{cases} + x & \text{if } x < \text{lower} \\ + x & \text{if } x > \text{upper} \\ + 0 & \text{otherwise} + \end{cases}. + $$ + + Usage: + + >>> x = tf.constant([1.0, 0.0, 1.0]) + >>> kca.activations.hardshrink(x) + + + Args: + x: A `Tensor`. Must be one of the following types: + `bfloat16`, `float16`, `float32`, `float64`. + lower: `float`, lower bound for setting values to zeros. + upper: `float`, upper bound for setting values to zeros. + Returns: + A `Tensor`. Has the same type as `x`. + """ + if lower > upper: + raise ValueError( + "The value of lower is {} and should not be higher than the value variable upper, which is {} .".format( + lower, upper) + ) + mask_lower = x < lower + mask_upper = upper < x + mask = ops.logical_or(mask_lower, mask_upper) + mask = ops.cast(mask, x.dtype) + return x * mask diff --git a/kca/activations/lisht.py b/kca/activations/lisht.py new file mode 100644 index 0000000..c208761 --- /dev/null +++ b/kca/activations/lisht.py @@ -0,0 +1,28 @@ +from keras_core import activations +from kca.utils.types import TensorLike + + +def lisht(x: TensorLike) -> TensorLike: + r"""LiSHT: Non-Parameteric Linearly Scaled Hyperbolic Tangent Activation Function. + + Computes linearly scaled hyperbolic tangent (LiSHT): + + $$ + \mathrm{lisht}(x) = x * \tanh(x). + $$ + + See [LiSHT: Non-Parameteric Linearly Scaled Hyperbolic Tangent Activation Function for Neural Networks](https://arxiv.org/abs/1901.05894). + + Usage: + + >>> x = tf.constant([1.0, 0.0, 1.0]) + >>> kca.activations.lisht(x) + + + Args: + x: A `Tensor`. Must be one of the following types: + `bfloat16`, `float16`, `float32`, `float64`. + Returns: + A `Tensor`. Has the same type as `x`. + """ + return x * activations.tanh(x) diff --git a/kca/activations/mish.py b/kca/activations/mish.py new file mode 100644 index 0000000..9f1dbb2 --- /dev/null +++ b/kca/activations/mish.py @@ -0,0 +1,28 @@ +from keras_core import activations +from kca.utils.types import TensorLike + + +def mish(x: TensorLike) -> TensorLike: + r"""Mish: A Self Regularized Non-Monotonic Neural Activation Function. + + Computes mish activation: + + $$ + \mathrm{mish}(x) = x \cdot \tanh(\mathrm{softplus}(x)). + $$ + + See [Mish: A Self Regularized Non-Monotonic Neural Activation Function](https://arxiv.org/abs/1908.08681). + + Usage: + + >>> x = tf.constant([1.0, 0.0, 1.0]) + >>> kca.activations.mish(x) + + + Args: + x: A `Tensor`. Must be one of the following types: + `bfloat16`, `float16`, `float32`, `float64`. + Returns: + A `Tensor`. Has the same type as `x`. + """ + return x * activations.tanh(activations.softplus(x)) diff --git a/kca/activations/rrelu.py b/kca/activations/rrelu.py new file mode 100644 index 0000000..20d0e3b --- /dev/null +++ b/kca/activations/rrelu.py @@ -0,0 +1,75 @@ +from keras_core import ops, random +from typing import Optional +from kca.utils.types import TensorLike, Number, Generator + + +def rrelu(x: TensorLike, lower: Number = 0.125, upper: Number = 0.3333333333333333, training: bool = False, + seed: Optional[int] = None, rng: Optional[Generator] = None) -> TensorLike: + r"""Randomized leaky rectified liner unit function. + + Computes rrelu function: + + $$ + \mathrm{rrelu}(x) = + \begin{cases} + x & \text{if } x > 0 \\ + a x + \end{cases}, + $$ + + where + + $$ + a \sim \mathcal{U}(\mathrm{lower}, \mathrm{upper}) + $$ + + when `training` is `True`; or + + $$ + a = \frac{\mathrm{lower} + \mathrm{upper}}{2} + $$ + + when `training` is `False`. + + See [Empirical Evaluation of Rectified Activations in Convolutional Network](https://arxiv.org/abs/1505.00853). + + Usage: + + >>> x = tf.constant([-1.0, 0.0, 1.0]) + >>> kca.activations.rrelu(x, training=False) + + >>> kca.activations.rrelu(x, training=True, seed=2020) + + >>> generator = tf.random.Generator.from_seed(2021) + >>> kca.activations.rrelu(x, training=True, rng=generator) + + + Args: + x: A `Tensor`. Must be one of the following types: + `bfloat16`, `float16`, `float32`, `float64`. + lower: `float`, lower bound for random alpha. + upper: `float`, upper bound for random alpha. + training: `bool`, indicating whether the `call` + is meant for training or inference. + seed: `int`, this sets the operation-level seed. + rng: A `tf.random.Generator`. + Returns: + result: A `Tensor`. Has the same type as `x`. + """ + lower = ops.cast(lower, x.dtype) + upper = ops.cast(upper, x.dtype) + + def random_a(): + if rng is not None and seed is not None: + raise ValueError("Either seed or rng should be specified. Not both at the same time.") + + if rng is not None: + return rng.uniform(ops.shape(x), minval=lower, maxval=upper, dtype=x.dtype) + + return random.uniform( + ops.shape(x), minval=lower, maxval=upper, dtype=x.dtype, seed=seed + ) + + a = random_a() if training else ((lower + upper) / 2) + + return ops.where(x >= 0, x, a * x) diff --git a/kca/activations/snake.py b/kca/activations/snake.py new file mode 100644 index 0000000..d15ef20 --- /dev/null +++ b/kca/activations/snake.py @@ -0,0 +1,30 @@ +from keras_core import ops, activations +from kca.utils.types import TensorLike, Number + + +def snake(x: TensorLike, frequency: Number = 1) -> TensorLike: + r"""Snake activation to learn periodic functions. + + Computes snake activation: + + $$ + \mathrm{snake}(x) = \mathrm{x} + \frac{1 - \cos(2 \cdot \mathrm{frequency} \cdot x)}{2 \cdot \mathrm{frequency}}. + $$ + + See [Neural Networks Fail to Learn Periodic Functions and How to Fix It](https://arxiv.org/abs/2006.08195). + + Usage: + + >>> x = tf.constant([-1.0, 0.0, 1.0]) + >>> tfa.activations.snake(x) + + + Args: + x: A `Tensor`. + frequency: A scalar, frequency of the periodic part. + Returns: + A `Tensor`. Has the same type as `x`. + """ + frequency = ops.cast(frequency, x.dtype) + + return x + (1 - ops.cos(2 * frequency * x)) / (2 * frequency) diff --git a/kca/activations/softshrink.py b/kca/activations/softshrink.py new file mode 100644 index 0000000..3dd38a6 --- /dev/null +++ b/kca/activations/softshrink.py @@ -0,0 +1,41 @@ +from keras_core import ops + +from kca.utils.types import TensorLike, Number + + +def softshrink(x: TensorLike, lower: Number = -0.5, upper: Number = 0.5) -> TensorLike: + r"""Soft shrink function. + + Computes soft shrink function: + + $$ + \mathrm{softshrink}(x) = + \begin{cases} + x - \mathrm{lower} & \text{if } x < \mathrm{lower} \\ + x - \mathrm{upper} & \text{if } x > \mathrm{upper} \\ + 0 & \text{otherwise} + \end{cases}. + $$ + + Usage: + + >>> x = tf.constant([-1.0, 0.0, 1.0]) + >>> kca.activations.softshrink(x) + + + Args: + x: A `Tensor`. Must be one of the following types: + `bfloat16`, `float16`, `float32`, `float64`. + lower: `float`, lower bound for setting values to zeros. + upper: `float`, upper bound for setting values to zeros. + Returns: + A `Tensor`. Has the same type as `x`. + """ + if lower > upper: + raise ValueError( + "The value of lower is {} and should not be higher than the value variable upper, which is {} .".format( + lower, upper) + ) + values_below_lower = ops.where(x < lower, x - lower, 0) + values_above_upper = ops.where(upper < x, x - upper, 0) + return values_below_lower + values_above_upper diff --git a/kca/activations/tanhshrink.py b/kca/activations/tanhshrink.py new file mode 100644 index 0000000..6b9cc09 --- /dev/null +++ b/kca/activations/tanhshrink.py @@ -0,0 +1,26 @@ +from keras_core import activations +from kca.utils.types import TensorLike + + +def tanhshrink(x: TensorLike) -> TensorLike: + r"""Tanh shrink function. + + Applies the element-wise function: + + $$ + \mathrm{tanhshrink}(x) = x - \tanh(x). + $$ + + Usage: + + >>> x = tf.constant([-1.0, 0.0, 1.0]) + >>> kca.activations.tanhshrink(x) + + + Args: + x: A `Tensor`. Must be one of the following types: + `bfloat16`, `float16`, `float32`, `float64`. + Returns: + A `Tensor`. Has the same type as `x`. + """ + return x - activations.tanh(x) diff --git a/kca/utils/types.py b/kca/utils/types.py index 7ad52cb..acec805 100644 --- a/kca/utils/types.py +++ b/kca/utils/types.py @@ -61,4 +61,8 @@ OptimizerType = Union[ keras.optimizers.Optimizer, str -] \ No newline at end of file +] + +Generator = Union[ + tf.random.Generator +]