Introduction
Deep learning has revolutionized the field of artificial intelligence, and at the heart of this revolution are Convolutional Neural Networks (CNNs). CNNs have become the go-to architectures for tasks involving images, such as object detection, facial recognition, medical imaging, and self-driving cars.
Traditionally, frameworks like TensorFlow and PyTorch have dominated the deep learning landscape. However, JAX has emerged as a powerful alternative, especially for research and high-performance computing. Developed by Google, JAX provides automatic differentiation and Just-In-Time (JIT) compilation, making it highly efficient for numerical computing and deep learning applications.
Why Use JAX for CNNs?
JAX stands out because of its ability to seamlessly run code on CPUs, GPUs, and TPUs while maintaining a NumPy-like API. This means you can develop models using familiar syntax while benefiting from:
Automatic Vectorization – With functions like vmap, JAX makes it easy to apply operations over large batches of data without writing explicit loops.
Efficient Autograd – JAX provides automatic differentiation using grad, which simplifies training deep learning models.
XLA Compilation – Just-In-Time (JIT) compilation speeds up execution by compiling computation graphs for efficient hardware utilization.
Functional Programming Paradigm – Unlike traditional deep learning frameworks, JAX encourages pure functions, which improves reproducibility and debugging.
Prerequisites
Before proceeding, ensure you are familiar with:
JAX fundamentals, you can check out the JAX documentation here.
Building CNNs in TensorFlow or PyTorch
JAX optimizers and loss functions
Install Dependencies
Install the required libraries:
!pip install jax jaxlib flax optax tensorflow tensorflow_datasets dm-pix tqdm matplotlib
jax and jaxlib – The core JAX library and its hardware acceleration backend.
flax – A neural network library for JAX, similar to PyTorch’s torch.nn.
optax – A library for optimization algorithms in JAX.
dm_pix – A lightweight image processing library for JAX.
matplotlib – For visualizing images.
Import Packages
Load necessary libraries:
import os
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt
import jax
import jax.numpy as jnp
import optax
from tqdm.auto import tqdm
from flax import linen as nn
from flax.training import train_state
import dm_pix as pix
Verify GPU Access
JAX runs computations on CPUs, GPUs, and TPUs seamlessly. To check if your machine has a GPU or TPU available:
print("Available Devices:", jax.devices())
if jax.default_backend() == "gpu":
print("Using GPU:", jax.devices("gpu"))
elif jax.default_backend() == "tpu":
print("Using TPU:", jax.devices("tpu"))
else:
print("Using CPU")
Data Preprocessing
Raw image data comes in various sizes, orientations, and quality levels. Preprocessing is crucial for:
Ensuring uniform input dimensions.
Normalizing pixel values for stable training.
Augmenting data to improve model generalization.
Converting images into JAX-compatible tensors.
Loading the Dataset
We will use the Cats vs. Dogs dataset available on Kaggle. Download and unzip the dataset using:
!kaggle datasets download -d chetankv/dogs-cats-images
!unzip dogs-cats-images.zip
This dataset gives us:
Define the path to the images and the batch size:
base_dir = "/content/dog vs cat/dataset/training_set"
batch_size = 64
Resizing and Normalizing Images
Since images in the dataset have varying sizes, we must resize them to a fixed size (e.g., 128×128 pixels). Additionally, we normalize pixel values from [0, 255] → [0, 1] for stable training.
IMG_SIZE = 128
resize_and_rescale = tf.keras.Sequential(
[
tf.keras.layers.Resizing(IMG_SIZE, IMG_SIZE),
tf.keras.layers.Rescaling(1.0 / 255),
]
)
Data Augmentation for Better Generalization
Data augmentation helps improve model generalization by applying transformations like flipping, rotation, brightness adjustments, and cropping.
rng = jax.random.PRNGKey(0)
rng, inp_rng, init_rng = jax.random.split(rng, 3)
delta = 0.42
factor = 0.42
@jax.jit
def data_augmentation(image):
new_image = pix.adjust_brightness(image=image, delta=delta)
new_image = pix.random_brightness(image=new_image, max_delta=delta, key=inp_rng)
new_image = pix.flip_up_down(image=image)
new_image = pix.flip_left_right(image=new_image)
new_image = pix.rot90(k=1, image=new_image)
return new_image
Converting Data to JAX-Compatible Tensors
JAX primarily operates on NumPy-like arrays (jnp.array). TensorFlow uses tf.Tensor, so we must convert our dataset into a JAX-friendly format.
AUTOTUNE = tf.data.AUTOTUNE
def prepare(ds, shuffle=False):
ds = ds.map(lambda x, y: (resize_and_rescale(x), y), num_parallel_calls=AUTOTUNE)
if shuffle:
ds = ds.shuffle(1000)
return ds.prefetch(buffer_size=AUTOTUNE)
train_ds = prepare(training_set, shuffle=True)
val_ds = prepare(validation_set)
evaluation_set = prepare(eval_set)
Visualizing Preprocessed Images
Let’s check if preprocessing is working as expected.
plt.figure(figsize=(10, 10))
augmented_images = []
for images, _ in training_set.take(1):
for i in range(9):
augmented_image = data_augmentation(np.array(images[i], dtype=jnp.float32))
augmented_images.append(augmented_image)
ax = plt.subplot(3, 3, i + 1)
plt.imshow(augmented_images[i].astype("uint8"))
plt.axis("off")
Defining a CNN in JAX
In JAX, neural networks are often implemented using Flax, a high-level neural network library designed to work seamlessly with JAX’s functional paradigm. Flax provides an intuitive way to define models using Module classes.
Below is a simple implementation of a Convolutional Neural Network (CNN) in JAX using Flax:
import jax.numpy as jnp
import flax.linen as nn
class CNN(nn.Module):
num_classes: int
@nn.compact
def __call__(self, x):
x = nn.Conv(features=32, kernel_size=(3, 3), strides=(1, 1), padding="SAME")(x)
x = nn.relu(x)
x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
x = nn.Conv(features=64, kernel_size=(3, 3), strides=(1, 1), padding="SAME")(x)
x = nn.relu(x)
x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
x = nn.Conv(features=128, kernel_size=(3, 3), strides=(1, 1), padding="SAME")(x)
x = nn.relu(x)
x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
x = x.reshape((x.shape[0], -1))
x = nn.Dense(features=128)(x)
x = nn.relu(x)
x = nn.Dense(features=self.num_classes)(x)
return x
First Convolutional Layer
Applies a 32-channel convolution with a 3×3 kernel.
Uses ReLU activation to introduce non-linearity.
Applies max pooling with a 2×2 window and a stride of 2, reducing the spatial dimensions by half.
Second Convolutional Layer
Uses a 64-channel convolution with a 3×3 kernel.
Again applies ReLU activation.
Another max pooling operation further reduces spatial dimensions.
Third Convolutional Layer
Flattening and Fully Connected Layers
The feature maps from the final convolutional layer are flattened into a 1D vector.
A dense layer with 128 neurons applies a ReLU activation.
The final output layer produces logits corresponding to the number of classes.
Why Use @nn.compact?
Flax provides two ways to define models:
Using @nn.compact, which allows direct instantiation of layers within the __call__ method.
Using setup(), where layers are explicitly defined as attributes.
The compact approach is cleaner and more intuitive for simple models, avoiding the need to define layer attributes separately.
Initializing the Model
JAX does not use an implicit state, so model parameters must be explicitly initialized. The init function from Flax helps generate the model’s parameters using a random key and an input shape.
from flax.core import freeze, unfreeze
key = jax.random.PRNGKey(0)
input_shape = (1, 32, 32, 3)
model = CNN(num_classes=10)
params = model.init(key, jnp.ones(input_shape))["params"]
A random key is generated using jax.random.PRNGKey(0). JAX requires explicit control over random number generation for reproducibility.
A dummy input tensor of shape (1, 32, 32, 3) is created to initialize the network.
The model is instantiated and the init function generates model parameters using the random key.
The "params" field is extracted from the initialization output, as Flax’s init method returns a dictionary containing additional information (e.g., batch statistics if using BatchNorm).
Defining the Training State
Flax provides a train_state abstraction to manage model parameters, optimizer state, and other training-related information. The optax library is used for defining the optimizer.
import optax
from flax.training import train_state
class TrainState(train_state.TrainState):
pass
learning_rate = 0.001
optimizer = optax.adam(learning_rate)
state = TrainState.create(apply_fn=model.apply, params=params, tx=optimizer)
TrainState is a dataclass that stores the model's parameters, optimizer state, and apply_fn (the function used for forward passes).
Optax's Adam optimizer is set up with a learning rate of 0.001.
The state.create() method initializes the model’s training state with:
apply_fn: The forward pass function from the model.
params: The initialized parameters from the previous step.
tx: The optimizer (Adam in this case).
Defining Loss and Accuracy Metrics
A loss function is required to guide training, while an accuracy function evaluates model performance.
Loss Function
import jax.nn as jnn
def cross_entropy_loss(params, state, batch):
logits = state.apply_fn({'params': params}, batch['images'])
labels = jnn.one_hot(batch['labels'], num_classes=10)
return -jnp.sum(labels * jnn.log_softmax(logits)) / batch['labels'].shape[0]
The function takes model parameters, the current training state, and a batch of input data.
The forward pass is performed using apply_fn, producing logits (raw model predictions).
The labels are one-hot encoded to match the logits' shape.
The cross-entropy loss is computed using log_softmax(logits), ensuring numerical stability.
The loss is averaged over the batch size for proper optimization.
Accuracy Function
def compute_accuracy(params, state, batch):
logits = state.apply_fn({'params': params}, batch['images'])
predicted_labels = jnp.argmax(logits, axis=1)
return jnp.mean(predicted_labels == batch['labels'])
The function takes model parameters, the training state, and a batch of data.
The forward pass is executed to obtain logits.
The highest-scoring class is selected using argmax(), determining the model’s predicted label.
Accuracy is computed by comparing predictions with actual labels and averaging the correct classifications.
Training and Evaluating a CNN in JAX
Training in JAX is based on functional transformations, meaning explicit gradient computation and parameter updates are required. The jax.grad function is used to compute gradients efficiently.
Training Step Function
@jax.jit
def train_step(state, batch):
def loss_fn(params):
logits = state.apply_fn({'params': params}, batch['images'])
labels = jax.nn.one_hot(batch['labels'], num_classes=10)
loss = -jnp.sum(labels * jax.nn.log_softmax(logits)) / batch['labels'].shape[0]
return loss
grads = jax.grad(loss_fn)(state.params)
state = state.apply_gradients(grads=grads)
return state
JIT Compilation (@jax.jit): JAX’s just-in-time compilation speeds up training by optimizing computation.
loss_fn Function: Defines the cross-entropy loss to be minimized.
jax.grad(loss_fn): Computes gradients with respect to model parameters.
state.apply_gradients(grads=grads): Updates the training state using computed gradients.
Training Loop
def train_model(state, train_loader, num_epochs=10):
for epoch in range(num_epochs):
for batch in train_loader:
state = train_step(state, batch)
print(f"Epoch {epoch + 1} completed")
return state
Iterates through multiple epochs, training the model for num_epochs.
Processes each batch, updating the model parameters.
Logs progress at the end of each epoch.
Evaluation Step Function
def evaluate_model(state, test_loader):
accuracies = []
for batch in test_loader:
acc = compute_accuracy(state.params, state, batch)
accuracies.append(acc)
final_accuracy = jnp.mean(jnp.array(accuracies))
print(f"Test Accuracy: {final_accuracy * 100:.2f}%")
return final_accuracy
Iterates through the test dataset, computing accuracy for each batch.
Aggregates accuracy scores across batches to compute the final accuracy.
Prints the test accuracy, indicating how well the model generalizes to unseen data.
Conclusion
JAX’s functional and hardware-accelerated approach allows for efficient model training, particularly on GPUs and TPUs. The explicit handling of gradients and optimizers ensures flexibility while maintaining high performance.
Future work could explore advanced techniques such as data augmentation, model regularization, and hyperparameter tuning to improve performance. Additionally, integrating JAX with frameworks like TensorFlow or PyTorch could provide hybrid workflows for deep learning research and production applications.
