Creating Custom Activation Functions with Lambda Layers in TensorFlow 2 | by Arjun Sarkar | Towards Data Science

Previously we’ve seen how to create custom loss functions — Creating custom Loss functions using TensorFlow 2
Introduction
In this article, we look at how to create custom activation functions. While TensorFlow already contains a bunch of activation functions inbuilt, there are ways to create your own custom activation function or to edit an existing activation function.
ReLU (Rectified Linear Unit) is still the most common activation function used in the hidden layers of any neural network architecture. ReLU can also be represented as a function f(x) where,
f(x) = 0, when x <0,
and, f(x) = x, when x ≥ 0.
Thus the function takes into consideration only the positive part, and is written as,
f(x) = max(0,x)
or in a code representation,
if input > 0:   return inputelse:   return 0
But this ReLU function is predefined. What if we want to customize this function or create our own ReLU activation. There is a very simple way to do this in TensorFlow — we just have to use Lambda layers.
How to use a lambda layer?
tf.keras.layers.Lambda(lambda x: tf.abs(x))
Lambda is just another layer that can be called directly in TensorFlow. Within the lambda layer, first, the parameter is specified. In the above code snippet, this value is ‘x’ (lambda x). In this case, we want to find the absolute value of x, so we use tf.abs(x). So if x has a value of -1, this lambda layer changes the value of x to 1.
How to use a lambda layer to create a custom ReLU?
def custom_relu(x):    return K.maximum(0.0,x)model = tf.keras.models.Sequential([     tf.keras.layers.Flatten(input_shape=(128,128)),     tf.keras.layers.Dense(512),     tf.keras.layers.Lambda(custom_relu),     tf.keras.layers.Dense(5, activation = 'softmax')])
The above code snippet shows how a custom ReLU can be implemented in a TensorFlow model. We create a function custom_relu and return the maximum of 0 or x (same as the ReLU function would do).
In the sequential model below, after the Dense layer, we create a Lambda layer and pass it in the custom activation function. But this code still does not do anything different than what a ReLU activation function would do.
The fun starts when we start playing around with the return value of our custom function. Let's say we take the maximum of 0.5 and x, rather than 0 and x. There we have out, our own custom ReLU. These values can then be altered as required.
def custom_relu(x): return K.maximum(0.5,x)
def custom_relu(x):    return K.maximum(0.5,x)model = tf.keras.models.Sequential([     tf.keras.layers.Flatten(input_shape=(128,128)),     tf.keras.layers.Dense(512),     tf.keras.layers.Lambda(custom_relu),     tf.keras.layers.Dense(5, activation = 'softmax')])
Example of using lambda activations on the mnist dataset
#using absolute value (Lambda layer example 1)import tensorflow as tffrom tensorflow.keras import backend as Kmnist = tf.keras.datasets.mnist(x_train, y_train),(x_test, y_test) = mnist.load_data()x_train, x_test = x_train / 255.0, x_test / 255.0model = tf.keras.models.Sequential([  tf.keras.layers.Flatten(input_shape=(28, 28)),  tf.keras.layers.Dense(128),  tf.keras.layers.Lambda(lambda x: tf.abs(x)),   tf.keras.layers.Dense(10, activation='softmax')])model.compile(optimizer='adam',              loss='sparse_categorical_crossentropy',              metrics=['accuracy'])model.fit(x_train, y_train, epochs=5)model.evaluate(x_test, y_test)
Replacing the ReLU activation with the absolute values on the mnist dataset gives a test accuracy of 97.384%.
#using custom ReLU activation (Lambda layer example 2)import tensorflow as tffrom tensorflow.keras import backend as Kmnist = tf.keras.datasets.mnist(x_train, y_train),(x_test, y_test) = mnist.load_data()x_train, x_test = x_train / 255.0, x_test / 255.0def my_relu(x):    return K.maximum(-0.1, x)model = tf.keras.models.Sequential([    tf.keras.layers.Flatten(input_shape=(28, 28)),    tf.keras.layers.Dense(128),    tf.keras.layers.Lambda(my_relu),     tf.keras.layers.Dense(10, activation='softmax')])model.compile(optimizer='adam',              loss='sparse_categorical_crossentropy',              metrics=['accuracy'])model.fit(x_train, y_train, epochs=5)model.evaluate(x_test, y_test)
Replacing the ReLU activation with the custom ReLU activation, taking a maximum of -0.1 or x, on the mnist dataset gives a test accuracy of 97.778%.
Conclusion
Even though lambda layers are very simple to use, they have many limitations. In the next article, I will write about creating fully custom layers in TensorFlow that are also trainable.