TensorFlow
Python
custom activation function
deep learning
machine learning

How to make a custom activation function with only Python in Tensorflow?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

In machine learning, particularly deep learning, activation functions play a vital role by introducing non-linearity into neural networks. While TensorFlow provides a suite of pre-built activation functions like ReLU, Sigmoid, and Tanh, there might be scenarios where specific tasks require a custom activation function. This guide will cover the process of creating a custom activation function using Python in TensorFlow.

Understanding Activation Functions

Activation functions determine the output of neurons in a network and define the output shape given inputs. They help models learn complex patterns and are essential for allowing networks to generalize.

Common Activation Functions:

  • ReLU (Rectified Linear Unit): f(x)=max(0,x)f(x) = \max(0, x)
  • Sigmoid: f(x)=11+exf(x) = \frac{1}{1 + e^{-x}}
  • Tanh: f(x)=exexex+exf(x) = \frac{e^x - e^{-x}}{e^x + e^{-x}}

These functions fit standard needs, but customizing them allows you to tailor behavior to specific problems, improving model performance.

Creating a Custom Activation Function

We'll walk through the steps of creating a simple custom activation function, called "Swish", defined as f(x)=xsigmoid(βx)f(x) = x \cdot \text{sigmoid}(\beta x), where β\beta is a parameter.

Step 1: Import Necessary Libraries

You first need to ensure you have TensorFlow installed. You can do this using pip:

bash
pip install tensorflow

Next, import the required libraries:

python
import tensorflow as tf
from tensorflow.keras.layers import Activation
from tensorflow.keras.utils import get_custom_objects

Step 2: Define the Activation Function

Define the function using standard Python syntax. Ensure it is compatible with TensorFlow operations.

python
def swish_activation(x, beta=1.0):
    return x * tf.keras.backend.sigmoid(beta * x)

Here, swish_activation takes an input tensor x and an optional parameter beta to control the shape of the function.

Step 3: Register the Custom Activation Function

To ensure TensorFlow recognizes the custom function, use get_custom_objects to register it.

python
get_custom_objects().update({'swish': Activation(swish_activation)})

Step 4: Use the Custom Activation Function in a Model

Now that you have defined and registered the custom function, you can use it within a Keras model like any built-in activation.

python
1model = tf.keras.models.Sequential([
2    tf.keras.layers.Dense(128, input_dim=784),
3    Activation(swish_activation),
4    tf.keras.layers.Dense(10, activation='softmax')
5])

Technical Explanation

The custom activation function leverages TensorFlow's backend operations to ensure compatibility with its computational graph. By expressing the function mathematically within TensorFlow's operations, it benefits from performance optimizations and can be used in both training and inference without any change.

Summary Table

Activation FunctionDefinitionCharacteristicsUsage
ReLUmax(0,x)\max(0, x)Non-linear, simple to computeGeneral-purpose
Sigmoid$\frac{1}{1 + e^{-x}}$Squashes input to $[0, 1]$Binary problems
Tanh$\frac{e^x - e^{-x}}{e^x + e^{-x}}$Squashes input to $[-1, 1]$Regression problems
Swish*xsigmoid(βx)x \cdot \text{sigmoid}(\beta x)Smooth, self-gatedCustom use-cases

Note: *Swish is a custom function created in this article.

Alternative Approaches and Considerations

  • Parameter Tuning: Experiment with different values of beta to control the non-linearity of the Swish function.
  • Graph Mode vs. Eager Execution: TensorFlow 2.x's eager execution mode makes debugging easier and is generally enabled by default. Custom functions will work seamlessly under this mode.
  • Performance Optimization: Pre-build the custom function within TensorFlow's tf.function decorator to leverage graph execution benefits.

Conclusion

Creating a custom activation function in TensorFlow can be beneficial for specific tasks that require non-standard activation behavior. By defining and integrating custom functions, as shown with the Swish activation, you can expand TensorFlow's versatility and tailor networks to specific needs. Experimentation and adaptation of these functions can often lead to improved performance and better generalization of models.


Course illustration
Course illustration

All Rights Reserved.