Search⌘ K
AI Features
Home
Courses
Deep Learning with JAX and Flax

Activation Functions in JAX

Explore common activation functions in JAX like ReLU, Sigmoid, Softmax, and advanced options such as ELU and GELU. Understand their roles in neural networks for different tasks and learn to define custom functions like LeakyReLU for tailored performance.

We'll cover the following...

Overview

Activation functions are applied in neural networks to ensure the network outputs the desired result. The activation function caps the output within a specific range. For instance, when solving a binary classification problem, the outcome should be a number between 0 and 1. This indicates the probability of an item belonging to either of the two classes.

Neural network
Neural network

However, in a regression problem, we want the numerical prediction of a quantity, for example, the price of an item. We should, therefore, choose an appropriate activation function for the problem being solved. Let’s look at common activation functions in JAX and Flax.

ReLU

The Rectified Linear Unit (ReLU) activation function is primarily used in the hidden layers of neural networks to ensure non-linearity. The function caps all outputs to zero and above. Outputs below zero are returned as zero, while numbers above zero are returned as they are. This ensures that there are no negative numbers in the network.

Graph of the ReLU activation function
Graph of the ReLU activation function

Let’s understand how to apply the ReLU activation function in the following code snippet:

Python 3.8
import flax
from flax import linen as nn
class CNN(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Conv(features=32, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = nn.Conv(features=64, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))
        x = nn.Dense(features=256)(x)
        x = nn.relu(x)
        x = nn.Dense(features=2)(x)
        x = nn.log_softmax(x)
        return x

In lines 7, 10, and 14, we apply the ReLU activation ...