Week 1 taught us that JAX is fast. Week 2 showed us how to eliminate loops with vmap and compute gradients with grad. Week 3 gave us Flax NNX—a way to define neural networks that feels like PyTorch but runs like JAX.

But we haven't actually trained anything yet.

Today, that changes. We're going to write a complete training loop from scratch: forward pass, loss calculation, gradient computation, parameter updates, and evaluation. No model.fit(). No magic. Just explicit, controllable, JIT-compiled code that you own completely.

By the end of this article, you'll have trained a CNN on MNIST and understood every single line of code that made it happen.

Why Write Your Own Training Loop?

If you've used Keras, you know how easy training can be:

model.fit(x_train, y_train, epochs=10)

One line. Done. So why would anyone write hundreds of lines to do the same thing?

Because model.fit() is a black box. It works until it doesn't. And when it doesn't, when your loss explodes, when you need gradient accumulation, when you're debugging a custom layer, when you need to log something specific every 37 steps, you're stuck.

The training loop is where machine learning actually happens. If you don't control it, you don't control your model.

In production systems, you almost always need:

• Custom metrics that aren't built into the framework

• Gradient clipping to prevent exploding gradients

• Learning rate schedules that change based on validation loss

• Mixed precision training for speed

• Gradient accumulation when batches don't fit in memory

• Early stopping based on custom criteria

Writing your own loop means you can implement any of these with a few lines of code.

The Four Beats of Training

Every training loop, regardless of framework, follows the same rhythm:

1. Forward Pass: Run input through the model to get predictions

2. Loss Calculation: Compare predictions to ground truth

3. Backward Pass: Compute gradients of the loss with respect to parameters

4. Parameter Update: Adjust parameters using the gradients

In PyTorch, this looks like:

optimizer.zero_grad()        # Clear old gradients
predictions = model(x)       # Forward pass
loss = loss_fn(predictions, y)  # Loss calculation
loss.backward()              # Backward pass
optimizer.step()             # Parameter update

PyTorch hides a lot here. loss.backward() traverses a hidden computational graph and secretly updates .grad attributes on every parameter.

JAX makes everything explicit.

The Optax Library

Optax is JAX's optimization library. It handles gradient transformations and parameter updates.

The key insight about Optax is that optimizers are composable. Instead of monolithic optimizer classes, Optax gives you building blocks that you chain together:

import optax

# A simple Adam optimizer
optimizer = optax.adam(learning_rate=0.001)

# Adam with gradient clipping
optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),  # First, clip gradients
    optax.adam(learning_rate=0.001)   # Then, apply Adam
)

# Adam with weight decay (AdamW)
optimizer = optax.adamw(learning_rate=0.001, weight_decay=0.01)

This composability is powerful. Need to add gradient clipping to your existing optimizer? Add one line. Need a custom transformation? Write a function that takes gradients and returns modified gradients.

Common Optax Optimizers

Common Gradient Transformations

The nnx.Optimizer Wrapper

Optax is a functional library, it works with raw JAX arrays. But our models are NNX objects with structured parameters. We need a bridge.

nnx.Optimizer is that bridge:

from flax import nnx
import optax

# Create the Optax optimizer
tx = optax.adam(learning_rate=0.001)

# Wrap it with nnx.Optimizer
optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)

The wrt parameter is critical. It stands for "with respect to" and tells the optimizer which variables to update. nnx.Param means "only update trainable parameters", not batch normalization statistics or other state.

Once wrapped, updating parameters is simple:

optimizer.update(model, grads)

This applies the gradients to the model's parameters in place. No need to manually extract parameters, apply updates, and put them back.

Computing Gradients with nnx.value_and_grad

In PyTorch, you call loss.backward() to compute gradients. In JAX, we use value_and_grad.

Here's the pattern:

def loss_fn(model, x, y):
    logits = model(x)
    loss = compute_loss(logits, y)
    return loss

# Get both the loss value and the gradients
loss, grads = nnx.value_and_grad(loss_fn)(model, x, y)

Critical rule: nnx.value_and_grad computes gradients with respect to the first argument of the function. That's why model must be the first parameter.

The grads object has the same structure as the model's parameters—it's a pytree where each leaf is the gradient for the corresponding parameter.

The Complete Training Step

Let's put it all together. Here's a complete, JIT-compiled training step:

@nnx.jit
def train_step(model, optimizer, x, y):
    """Execute one training step."""
    
    def loss_fn(model):
        logits = model(x)
        loss = optax.softmax_cross_entropy_with_integer_labels(
            logits=logits, 
            labels=y
        ).mean()
        return loss, logits
    
    # Compute loss and gradients
    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)
    
    # Update parameters
    optimizer.update(model, grads)
    
    # Compute accuracy for logging
    predictions = jnp.argmax(logits, axis=-1)
    accuracy = jnp.mean(predictions == y)
    
    return loss, accuracy

A few things to note:

1. @nnx.jit: The entire function is JIT-compiled. The first call traces and compiles; subsequent calls execute the compiled code.

2. has_aux=True: Our loss function returns (loss, logits). The has_aux=True flag tells value_and_grad that the function returns auxiliary data alongside the loss. Gradients are only computed for the first returned value.

3. In-place updates: optimizer.update(model, grads) modifies the model in place. NNX handles the functional-to-stateful conversion internally.

Training a CNN on MNIST

Let's train a real model on real data. We'll use the CNN from Week 3 and the MNIST dataset.

Setup and Data Loading

import jax
import jax.numpy as jnp
from flax import nnx
import optax
import tensorflow_datasets as tfds
import tensorflow as tf

# Disable TensorFlow GPU to avoid conflicts with JAX
tf.config.set_visible_devices([], 'GPU')

# Load MNIST
def load_mnist(batch_size=32):
    """Load and preprocess MNIST dataset."""
    
    def preprocess(sample):
        image = tf.cast(sample['image'], tf.float32) / 255.0
        label = sample['label']
        return {'image': image, 'label': label}
    
    train_ds = tfds.load('mnist', split='train')
    test_ds = tfds.load('mnist', split='test')
    
    train_ds = (train_ds
                .map(preprocess)
                .shuffle(1024)
                .batch(batch_size, drop_remainder=True)
                .prefetch(1))
    
    test_ds = (test_ds
               .map(preprocess)
               .batch(batch_size, drop_remainder=True)
               .prefetch(1))
    
    return train_ds, test_ds

train_ds, test_ds = load_mnist(batch_size=32)

Define the Model

We'll use a slightly simplified version of the CNN from Week 3:

from functools import partial

class CNN(nnx.Module):
    """Convolutional Neural Network for MNIST."""
    
    def __init__(self, *, rngs: nnx.Rngs):
        self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)
        self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)
        self.linear1 = nnx.Linear(3136, 256, rngs=rngs)
        self.linear2 = nnx.Linear(256, 10, rngs=rngs)
        self.pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))
    
    def __call__(self, x):
        x = nnx.relu(self.conv1(x))
        x = self.pool(x)
        x = nnx.relu(self.conv2(x))
        x = self.pool(x)
        x = x.reshape(x.shape[0], -1)  # Flatten
        x = nnx.relu(self.linear1(x))
        x = self.linear2(x)
        return x

Initialize Model and Optimizer

# Initialize model
model = CNN(rngs=nnx.Rngs(0))

# Test forward pass
dummy_input = jnp.ones((1, 28, 28, 1))
dummy_output = model(dummy_input)
print(f"Model output shape: {dummy_output.shape}")  # (1, 10)

# Initialize optimizer
learning_rate = 0.005
momentum = 0.9
tx = optax.adamw(learning_rate, momentum)
optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)

# Display model structure
nnx.display(model)

Define Training and Evaluation Steps

@nnx.jit
def train_step(model, optimizer, batch):
    """Single training step."""
    
    def loss_fn(model):
        logits = model(batch['image'])
        loss = optax.softmax_cross_entropy_with_integer_labels(
            logits=logits,
            labels=batch['label']
        ).mean()
        return loss, logits
    
    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)
    optimizer.update(model, grads)
    
    accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == batch['label'])
    return loss, accuracy

@nnx.jit
def eval_step(model, batch):
    """Single evaluation step."""
    logits = model(batch['image'])
    loss = optax.softmax_cross_entropy_with_integer_labels(
        logits=logits,
        labels=batch['label']
    ).mean()
    accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == batch['label'])
    return loss, accuracy

The Training Loop

# Training configuration
num_epochs = 5
train_steps_per_epoch = 1000
eval_steps = 200

print("Starting training...")
print("=" * 60)

for epoch in range(num_epochs):
    # Training
    model.train()  # Set model to training mode
    train_loss, train_acc = 0.0, 0.0
    
    for step, batch in enumerate(train_ds.as_numpy_iterator()):
        if step >= train_steps_per_epoch:
            break
        
        # Convert to JAX arrays
        batch = {k: jnp.array(v) for k, v in batch.items()}
        
        loss, acc = train_step(model, optimizer, batch)
        train_loss += loss
        train_acc += acc
    
    train_loss /= train_steps_per_epoch
    train_acc /= train_steps_per_epoch
    
    # Evaluation
    model.eval()  # Set model to evaluation mode
    eval_loss, eval_acc = 0.0, 0.0
    eval_batches = 0
    
    for step, batch in enumerate(test_ds.as_numpy_iterator()):
        if step >= eval_steps:
            break
        
        batch = {k: jnp.array(v) for k, v in batch.items()}
        
        loss, acc = eval_step(model, batch)
        eval_loss += loss
        eval_acc += acc
        eval_batches += 1
    
    eval_loss /= eval_batches
    eval_acc /= eval_batches
    
    print(f"Epoch {epoch + 1}/{num_epochs}")
    print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f}")
    print(f"Eval Loss:  {eval_loss:.4f} | Eval Acc:  {eval_acc:.4f}")

print("Training complete!")

Using nnx.MultiMetric for Cleaner Tracking

NNX provides a MultiMetric helper for accumulating metrics across batches:

# Define metrics
metrics = nnx.MultiMetric(
    accuracy=nnx.metrics.Accuracy(),
    loss=nnx.metrics.Average('loss'),
)

@nnx.jit
def train_step_with_metrics(model, optimizer, metrics, batch):
    def loss_fn(model):
        logits = model(batch['image'])
        loss = optax.softmax_cross_entropy_with_integer_labels(
            logits=logits, labels=batch['label']
        ).mean()
        return loss, logits
    
    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)
    optimizer.update(model, grads)
    
    # Update metrics in place
    metrics.update(loss=loss, logits=logits, labels=batch['label'])

# In training loop:
for batch in train_ds:
    train_step_with_metrics(model, optimizer, metrics, batch)

# Get aggregated metrics
results = metrics.compute()
print(f"Loss: {results['loss']:.4f}, Accuracy: {results['accuracy']:.4f}")

# Reset for next epoch
metrics.reset()

This is cleaner than manual accumulation and handles edge cases (like different batch sizes) correctly.

Inference: Using the Trained Model

After training, making predictions is straightforward:

model.eval()  # Ensure evaluation mode (deterministic)

@nnx.jit
def predict(model, images):
    logits = model(images)
    return jnp.argmax(logits, axis=-1)

# Get predictions
test_batch = next(test_ds.as_numpy_iterator())
predictions = predict(model, jnp.array(test_batch['image']))

print(f"Predictions: {predictions[:10]}")
print(f"Actual: {test_batch['label'][:10]}")

Exercises

1. Add gradient clipping: Modify the optimizer to include optax.clip_by_global_norm(1.0). Train the model and compare the loss curves.

2. Learning rate schedule: Implement a learning rate that decays over time using optax.warmup_cosine_decay_schedule. Check the Optax documentation for the parameters.

3. Track more metrics: Add precision and recall tracking alongside accuracy. You'll need to compute true positives, false positives, and false negatives.

4. Early stopping: Modify the training loop to stop if validation accuracy doesn't improve for 3 consecutive epochs.

Quick Reference

import optax
from flax import nnx

# Optimizer Setup
tx = optax.adam(learning_rate=0.001)
optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)

# Training Step Pattern
@nnx.jit
def train_step(model, optimizer, batch):
    def loss_fn(model):
        logits = model(batch['x'])
        loss = compute_loss(logits, batch['y'])
        return loss, logits
    
    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)
    optimizer.update(model, grads)
    return loss

# Evaluation Step Pattern
@nnx.jit  
def eval_step(model, batch):
    logits = model(batch['x'])
    loss = compute_loss(logits, batch['y'])
    return loss

# Training/Eval Mode
model.train()  # Enable dropout, stochastic behavior
model.eval()   # Deterministic inference

What's Next

We've trained a model. But what happens when training goes wrong? What if loss becomes NaN? What if shapes don't match? What if the model silently produces garbage?

Next week, we dive into Optax in more depth: learning rate schedules, gradient clipping, weight decay, and building custom optimizer stacks. We'll make our training more robust and production-ready.

But if you've been following along, you might have noticed something uncomfortable: we've been passing arrays around manually, tracking parameters in tuples, and writing functions that return updated state. It works, but it doesn't feel like building neural networks. It feels like accounting.

If you've used PyTorch, you know how natural it is to define a model as a class, call model(x), and let the framework handle the rest. That ergonomic experience is what made PyTorch the dominant framework for research.

Today, we get that experience in JAX.

Flax NNX is a neural network library that gives you PyTorch-style classes and methods while compiling down to JAX's XLA backend. You write object-oriented code. JAX runs functional code. NNX bridges the gap.

By the end of this article, we'll have:

1. Understood why NNX exists and what problem it solves

2. Built our first neural network using nnx.Module

3. Learned how NNX handles the critical issue of random number generation

4. Created a CNN for image classification—the same architecture used in real production systems

Let's build something real.

The State Problem

To understand why Flax NNX is necessary, we need to understand the fundamental tension between PyTorch and JAX.

How PyTorch Handles State

In PyTorch, a model is an object that contains its own parameters:

import torch.nn as nn

class PyTorchMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(784, 128)
        self.linear2 = nn.Linear(128, 10)
    
    def forward(self, x):
        x = torch.relu(self.linear1(x))
        return self.linear2(x)

model = PyTorchMLP()
output = model(input_data)  # Parameters are hidden inside

This is intuitive. The nn.Linear layers create their own weight matrices internally. You don't see them, you don't manage them, they just exist.

But this "hidden state" creates problems for JAX. When you call jax.jit on a function, JAX traces through it and compiles a computational graph. If that function secretly reads or writes to hidden variables, JAX can't see those operations, and it can't optimize them.

How Pure JAX Handles State

JAX demands pure functions: functions where the output depends only on the inputs, with no side effects. If you want to update parameters, you must pass them in and return them out:

def forward(params, x):
    x = jax.nn.relu(x @ params['w1'] + params['b1'])
    return x @ params['w2'] + params['b2']

def train_step(params, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    new_params = update_params(params, grads)
    return new_params, loss  # Must return the new state

# You're always passing params around
params, loss = train_step(params, x_batch, y_batch)
params, loss = train_step(params, x_batch, y_batch)
params, loss = train_step(params, x_batch, y_batch)

This is explicit and JIT-compatible, but it's tedious. For a model with dozens of layers, managing nested dictionaries of parameters becomes error-prone.

How Flax NNX Solves This

NNX gives you the PyTorch interface while secretly doing the JAX bookkeeping. You write classes with attributes. NNX intercepts those attributes, extracts the parameters when needed for JIT compilation, and updates them in place when you're done.

from flax import nnx

class NNX_MLP(nnx.Module):
    def __init__(self, rngs: nnx.Rngs):
        self.linear1 = nnx.Linear(784, 128, rngs=rngs)
        self.linear2 = nnx.Linear(128, 10, rngs=rngs)
    
    def __call__(self, x):
        x = nnx.relu(self.linear1(x))
        return self.linear2(x)

model = NNX_MLP(rngs=nnx.Rngs(0))
output = model(input_data)  # Looks just like PyTorch

The syntax is nearly identical to PyTorch. But under the hood, NNX can extract the parameters, pass them through JAX transformations, and put them back—all without you writing a single line of state management code.

Building Your First NNX Model

Let's build a simple multi-layer perceptron step by step.

The Basic Structure

Every NNX model inherits from nnx.Module:

from flax import nnx
import jax.numpy as jnp

class SimpleMLP(nnx.Module):
    def __init__(self, hidden_dim: int, output_dim: int, *, rngs: nnx.Rngs):
        # Define layers as attributes
        self.linear1 = nnx.Linear(784, hidden_dim, rngs=rngs)
        self.linear2 = nnx.Linear(hidden_dim, output_dim, rngs=rngs)
    
    def __call__(self, x):
        # Define forward pass
        x = self.linear1(x)
        x = nnx.relu(x)
        x = self.linear2(x)
        return x

Key differences from PyTorch:

1. No super().__init__(): NNX uses Python metaclasses, so you don't need to call the parent constructor.

2. __call__ instead of forward: In Python, __call__ makes an object callable. NNX uses this standard convention rather than PyTorch's custom forward method.

3. The rngs parameter: This is required. We'll explain why in the next section.

Instantiating the Model

# Create a random number generator
rngs = nnx.Rngs(42)  # Seed for reproducibility

# Create the model
model = SimpleMLP(hidden_dim=128, output_dim=10, rngs=rngs)

# Test with dummy data
x = jnp.ones((32, 784))  # Batch of 32, 784 features each
output = model(x)

print(f"Input shape:  {x.shape}")       # (32, 784)
print(f"Output shape: {output.shape}")  # (32, 10)

That's it. The model is ready to use.

The Randomness Requirement

You might be wondering: why do we need to pass rngs everywhere?

The Problem with Hidden Randomness

In NumPy or PyTorch, random number generation uses a global state:

import numpy as np
np.random.seed(42)
print(np.random.randn())  # Always the same value
print(np.random.randn())  # Different value (global state advanced)

This global state is invisible. It changes every time you call a random function. And as we established in Week 1, JAX cannot work with hidden, changing state.

JAX's Explicit Randomness

In JAX, randomness is deterministic. You create a "key," and that key always produces the same random numbers:

from jax import random

key = random.PRNGKey(42)
print(random.normal(key, shape=(3,)))  # Always the same
print(random.normal(key, shape=(3,)))  # Still the same!

To get different random numbers, you must split the key:

key, subkey = random.split(key)
print(random.normal(subkey, shape=(3,)))  # New values

How nnx.Rngs Helps

When you create a neural network, every layer needs random numbers to initialize its weights. If you have 50 layers, that's 50 key splits you'd need to manage manually.

nnx.Rngs automates this. It's a key dispenser that splits and distributes keys automatically:

rngs = nnx.Rngs(42)

# When you create a layer, it asks rngs for keys internally
linear1 = nnx.Linear(10, 20, rngs=rngs)  # Gets its own keys
linear2 = nnx.Linear(20, 30, rngs=rngs)  # Gets different keys

# Both layers have different, reproducible initializations

The critical benefit: reproducibility. If you and I both run nnx.Rngs(42), we get identical models. This matters for debugging, for scientific reproducibility, and for distributed training where multiple machines must initialize the same model.

Inspecting Your Model

PyTorch lets you print(model) to see the architecture. NNX has something better: nnx.display().

nnx.display(model)

This produces a rich, hierarchical view showing:

• Every layer and sublayer

• Parameter shapes and dtypes

• The total parameter count

• The structure of the computational graph

In Jupyter notebooks, this renders as an interactive tree you can expand and collapse. It's invaluable for debugging shape mismatches and verifying your architecture.

Building a CNN for Image Classification

Let's build something more substantial: a convolutional neural network for classifying images. This is the same architecture pattern used in production image classifiers.

The Architecture

We'll build a classic CNN with:

• Two convolutional blocks (Conv → ReLU → Pool)

• A flatten operation

• Two dense layers for classification

from flax import nnx
import jax.numpy as jnp
from functools import partial

class CNN(nnx.Module):
    """A simple CNN for image classification."""
    
    def __init__(self, num_classes: int, *, rngs: nnx.Rngs):
        # Convolutional layers
        self.conv1 = nnx.Conv(
            in_features=1,      # Input channels (1 for grayscale)
            out_features=32,    # Output channels
            kernel_size=(3, 3), # 3x3 filters
            rngs=rngs
        )
        self.conv2 = nnx.Conv(
            in_features=32,
            out_features=64,
            kernel_size=(3, 3),
            rngs=rngs
        )
        
        # Dense layers
        # After two 2x2 pooling operations on 28x28 input: 28 → 14 → 7
        # So we have 64 channels × 7 × 7 = 3136 features
        self.linear1 = nnx.Linear(3136, 256, rngs=rngs)
        self.linear2 = nnx.Linear(256, num_classes, rngs=rngs)
        
        # Pooling as a reusable operation
        self.pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))
    
    def __call__(self, x):
        # Block 1: Conv → ReLU → Pool
        x = self.conv1(x)
        x = nnx.relu(x)
        x = self.pool(x)
        
        # Block 2: Conv → ReLU → Pool
        x = self.conv2(x)
        x = nnx.relu(x)
        x = self.pool(x)
        
        # Flatten: (batch, height, width, channels) → (batch, features)
        x = x.reshape(x.shape[0], -1)
        
        # Classification head
        x = nnx.relu(self.linear1(x))
        x = self.linear2(x)
        
        return x

Testing the Model

# Initialize
model = CNN(num_classes=10, rngs=nnx.Rngs(0))

# Create dummy input: batch of 4 grayscale 28×28 images
# Shape: (batch, height, width, channels)
dummy_input = jnp.ones((4, 28, 28, 1))

# Forward pass
output = model(dummy_input)

print(f"Input shape:  {dummy_input.shape}")  # (4, 28, 28, 1)
print(f"Output shape: {output.shape}")       # (4, 10)

# Inspect the model structure
nnx.display(model)

The output shape is (4, 10)—four images, ten class scores each. This is exactly what we'd feed into a softmax for classification.

How NNX Compiles with JAX

Here's the magic: even though we're writing object-oriented code, we can still use JAX's transformations.

The @nnx.jit Decorator

NNX provides its own versions of JAX transforms that understand NNX objects:

@nnx.jit
def forward(model, x):
    return model(x)

# This is JIT-compiled, just like @jax.jit
output = forward(model, dummy_input)

When you call this function, NNX:

1. Extracts the model's parameters into a pure JAX pytree

2. Traces the computation with those parameters

3. Compiles the trace with XLA

4. Updates the model object with any changed state

You write familiar OOP code. JAX gets the pure functions it needs. Everyone wins.

Preview: The Split/Merge Pattern

Under the hood, NNX uses two key operations:

# Split: separate structure from state
graphdef, state = nnx.split(model)

# graphdef: the "blueprint" of the model (static)
# state: the actual parameter values (dynamic, a pytree)

# Merge: reconstruct the model from structure and state
reconstructed_model = nnx.merge(graphdef, state)

You rarely need to call these directly—@nnx.jit handles it automatically. But understanding this pattern helps when you need to do advanced things like:

• Saving and loading checkpoints (Week 9)

• Distributing models across devices (Week 10)

• Custom training loops with fine-grained control

Common Layers Reference

Here are the NNX equivalents of layers you know from PyTorch:

Activation functions aren't layers in NNX, they're just functions:

x = nnx.relu(x)
x = nnx.gelu(x)
x = nnx.softmax(x)
x = nnx.sigmoid(x)

Working with Parameters

Sometimes you need direct access to the parameters—for logging, for custom initialization, or for freezing layers.

Accessing Parameters

# Access a specific layer's weights
print(model.linear1.kernel.value.shape)  # (3136, 256)
print(model.linear1.bias.value.shape)    # (256,)

# Parameters are nnx.Param objects; .value gets the JAX array

Extracting All Parameters

# Get all parameters as a state object
state = nnx.state(model)

# Or specifically just the trainable parameters
params = nnx.state(model, nnx.Param)

Updating Parameters

# Update the model with new state
nnx.update(model, new_state)

This becomes important next week when we build training loops.

Exercises

Before moving on, try these:

1. Add Dropout: Modify the CNN to include nnx.Dropout(rate=0.5, rngs=rngs) between the dense layers. Note that dropout needs its own RNG stream for the random mask.

2. Build a deeper network: Create a 4-layer MLP with hidden dimensions [512, 256, 128, 64]. Use a loop in __init__ to avoid repetition.

3. Parameter counting: Write a function that takes an NNX model and returns the total number of trainable parameters. Hint: use nnx.state(model, nnx.Param) and jax.tree_util.tree_map.

Quick Reference

from flax import nnx
import jax.numpy as jnp

# Define a Model
class MyModel(nnx.Module):
    def __init__(self, *, rngs: nnx.Rngs):
        self.linear = nnx.Linear(10, 5, rngs=rngs)
    
    def __call__(self, x):
        return nnx.relu(self.linear(x))

# Instantiate
model = MyModel(rngs=nnx.Rngs(0))

# Forward Pass
output = model(jnp.ones((32, 10)))

# Inspect
nnx.display(model)

# Access Parameters
weights = model.linear.kernel.value
state = nnx.state(model, nnx.Param)

# JIT Compile
@nnx.jit
def forward(model, x):
    return model(x)

# Common Layers
nnx.Linear(in_features, out_features, rngs=rngs)
nnx.Conv(in_features, out_features, kernel_size, rngs=rngs)
nnx.BatchNorm(num_features, rngs=rngs)
nnx.Dropout(rate, rngs=rngs)
nnx.Embed(num_embeddings, features, rngs=rngs)

What's Next

We have a model. But a model that can't learn is just a random number generator with extra steps.

Next week, we build the training loop. We'll use:

• nnx.value_and_grad to compute loss and gradients

• nnx.Optimizer to manage parameter updates

• optax to define the optimization algorithm

• Metrics to track accuracy and loss

We'll train our CNN on real data and watch the loss curve drop. That's when this all becomes real.

But speed isn't JAX's only trick. The real power of JAX lies in its transformations; functions that take functions and return new functions with different behavior.

Today we're covering the two transformations that make JAX indispensable for machine learning:

1. jax.vmap: Automatic vectorization. Write code for one example, run it on a million.

2. jax.grad: Automatic differentiation. Get gradients of any function, for free.

By the end of this article, you'll understand why JAX developers almost never write for-loops, and you'll have trained multiple machine learning models in parallel without writing any loop at all.

The Problem with Python Loops

Let's start with why loops are the enemy.

When you write a Python for-loop, the interpreter does a surprising amount of work for each iteration:

results = []
for x in data:
    # For EACH iteration, Python must:
    # 1. Fetch x from memory
    # 2. Check the type of x
    # 3. Look up what "+" means for that type
    # 4. Execute the addition
    # 5. Append to the list (which may require memory reallocation)
    results.append(x + 1)

If data has a million elements, Python performs those administrative steps a million times. The actual math (x + 1) takes nanoseconds. The overhead takes microseconds. You're spending 99% of your time on bookkeeping.

In deep learning, we process batches: 64 images, 128 sentences, 256 audio clips at once. If you loop through them in Python, your GPU sits idle while Python shuffles paperwork.

NumPy helps by pushing operations into C:

results = data + 1  # Vectorized, fast

But what if your function is more complex than addition? What if it involves multiple steps, conditionals, or nested operations? You'd have to manually rewrite everything to handle batches, adding batch dimensions everywhere and keeping track of which axis is which.

This is where jax.vmap changes the game.

jax.vmap: Automatic Vectorization

vmap stands for "vectorizing map." It takes a function written for a single example and transforms it into a function that operates on batches.

The Basic Pattern

import jax
import jax.numpy as jnp

# A function that works on ONE number
def square(x):
    return x ** 2

# Transform it to work on MANY numbers
batched_square = jax.vmap(square)

# Now use it
numbers = jnp.array([1, 2, 3, 4, 5])
result = batched_square(numbers)
print(result)  # [1, 4, 9, 16, 25]

"But wait," you might say, "I could just write numbers ** 2 directly."

True. The power of vmap shows up with complex functions:

def complex_operation(x):
    """A function with multiple steps."""
    a = jnp.sin(x)
    b = jnp.exp(-x ** 2)
    c = a * b + jnp.log(1 + jnp.abs(x))
    return c

# Without vmap, you'd need to think about broadcasting at each step
# With vmap, you just wrap it
batched_complex = jax.vmap(complex_operation)

x_batch = jnp.linspace(-3, 3, 1000)
results = batched_complex(x_batch)

You write the function thinking about one input. vmap handles the batch dimension for you.

Multiple Arguments: The in_axes Parameter

Real functions have multiple arguments. in_axes tells vmap which arguments to map over and which to broadcast.

Consider a dot product:

def dot_product(weights, features):
    return jnp.dot(weights, features)

Different scenarios require different in_axes:

# Scenario 1: One set of weights, many feature vectors
# weights: don't map (None), features: map over axis 0
batch_predict = jax.vmap(dot_product, in_axes=(None, 0))

weights = jnp.array([1.0, 2.0, 3.0])
features = jnp.array([
    [1.0, 0.0, 0.0],
    [0.0, 1.0, 0.0],
    [0.0, 0.0, 1.0],
])

results = batch_predict(weights, features)
print(results)  # [1., 2., 3.]

# Scenario 2: Many weight sets, one feature vector (ensemble of models)
# weights: map over axis 0, features: don't map (None)
ensemble_predict = jax.vmap(dot_product, in_axes=(0, None))

many_weights = jnp.array([
    [1.0, 0.0, 0.0],
    [0.0, 1.0, 0.0],
    [0.0, 0.0, 1.0],
])
single_features = jnp.array([1.0, 2.0, 3.0])

results = ensemble_predict(many_weights, single_features)
print(results)  # [1., 2., 3.]

# Scenario 3: Many weights, many features (parallel evaluation)
# Both: map over axis 0
parallel_predict = jax.vmap(dot_product, in_axes=(0, 0))

results = parallel_predict(many_weights, features)
print(results)  # [1., 2., 3.]

The rule is simple:

• 0 means "iterate over the first axis of this argument"

• None means "broadcast this argument to all iterations"

• You can use other integers for different axes

Nested vmap

You can stack vmap calls for multi-dimensional batching:

def single_multiply(a, b):
    return a * b

# Map over rows, then over columns
double_batched = jax.vmap(jax.vmap(single_multiply))

matrix_a = jnp.array([[1, 2], [3, 4]])
matrix_b = jnp.array([[10, 20], [30, 40]])

result = double_batched(matrix_a, matrix_b)
print(result)
# [[10, 40],
#  [90, 160]]

jax.grad: Automatic Differentiation

The other transformation that makes JAX essential for ML is jax.grad. It computes gradients automatically.

The Basic Pattern

def f(x):
    return x ** 2

# grad returns a NEW FUNCTION that computes the derivative
df_dx = jax.grad(f)

print(f(3.0))      # 9.0
print(df_dx(3.0))  # 6.0 (derivative of x² is 2x, and 2*3=6)

This works for any function, no matter how complex:

def messy_function(x):
    return jnp.sin(x) * jnp.exp(-x ** 2) + jnp.tanh(x)

gradient_fn = jax.grad(messy_function)

# The gradient at x=1.0
print(gradient_fn(1.0))  # -0.5047...

You didn't write any derivative rules. JAX traced through your function and computed the gradient automatically using the chain rule.

jax.value_and_grad: Get Both at Once

In training loops, you need both the loss value (to log it) and the gradients (to update parameters). jax.value_and_grad gives you both in one pass:

def loss_fn(params, x, y):
    prediction = params * x
    return (prediction - y) ** 2

# Returns (loss_value, gradients)
loss_and_grad_fn = jax.value_and_grad(loss_fn)

params = 1.0
x, y = 2.0, 6.0  # We want params=3.0 so that 3*2=6

loss, grad = loss_and_grad_fn(params, x, y)
print(f"Loss: {loss}")  # 16.0 (because (1*2 - 6)² = 16)
print(f"Grad: {grad}")  # -16.0

Gradients with Respect to Specific Arguments

By default, grad differentiates with respect to the first argument. Use argnums to change this:

def f(x, y):
    return x ** 2 + x * y

# Gradient with respect to x (first argument, default)
df_dx = jax.grad(f, argnums=0)

# Gradient with respect to y (second argument)
df_dy = jax.grad(f, argnums=1)

# Gradients with respect to both
df_both = jax.grad(f, argnums=(0, 1))

x, y = 2.0, 3.0
print(f"df/dx: {df_dx(x, y)}")  # 2*2 + 3 = 7
print(f"df/dy: {df_dy(x, y)}")  # 2
print(f"Both:  {df_both(x, y)}")  # (7.0, 2.0)

Combining Transforms: The Real Power

JAX transforms compose. You can combine jit, vmap, and grad freely:

def loss_single(params, x, y):
    """Loss for a single data point."""
    pred = params[0] * x + params[1]  # Linear: y = mx + b
    return (pred - y) ** 2

# Stack the transforms:
# 1. grad: compute gradients with respect to params
# 2. vmap: do this for a batch of (x, y) pairs
# 3. jit: compile the whole thing

batched_grad_fn = jax.jit(
    jax.vmap(
        jax.grad(loss_single),
        in_axes=(None, 0, 0)  # Same params, batch of x, batch of y
    )
)

params = jnp.array([1.0, 0.0])  # Initial guess: y = 1*x + 0
x_batch = jnp.array([1.0, 2.0, 3.0])
y_batch = jnp.array([2.0, 4.0, 6.0])  # True relationship: y = 2x

# Get gradients for each example in the batch
grads_per_example = batched_grad_fn(params, x_batch, y_batch)
print("Gradients per example:")
print(grads_per_example)

# Average them for a batch gradient
batch_grad = jnp.mean(grads_per_example, axis=0)
print(f"Batch gradient: {batch_grad}")

This pattern, jit(vmap(grad(...))), is the backbone of efficient training in JAX.

Project: Parallel Linear Regression

Let's train multiple models simultaneously without a single Python loop during training.

Scenario: We have housing price data from three different cities. Each city has different price dynamics, so we want to train a separate linear regression model for each.

Step 1: Generate Synthetic Data

import jax
import jax.numpy as jnp
from jax import random

# Seed for reproducibility
key = random.PRNGKey(42)

# True parameters for 3 cities
# City 0: price = 50 * size + 100
# City 1: price = 30 * size + 200  
# City 2: price = 80 * size + 50
true_slopes = jnp.array([50.0, 30.0, 80.0])
true_intercepts = jnp.array([100.0, 200.0, 50.0])

n_cities = 3
n_samples = 100

# Generate features (house sizes) for each city
key, subkey = random.split(key)
X = random.uniform(subkey, (n_cities, n_samples, 1), minval=10, maxval=100)

# Generate targets (prices) with some noise
key, subkey = random.split(key)
noise = random.normal(subkey, (n_cities, n_samples)) * 50

# Y[i] = true_slopes[i] * X[i] + true_intercepts[i] + noise[i]
Y = (X[:, :, 0] * true_slopes[:, None] + 
     true_intercepts[:, None] + 
     noise)

print(f"X shape: {X.shape}")  # (3, 100, 1)
print(f"Y shape: {Y.shape}")  # (3, 100)

Step 2: Define the Model for ONE City

We write everything as if we only have one city:

def predict(params, x):
    """Predict price for one city's houses."""
    slope, intercept = params[0], params[1]
    return x[:, 0] * slope + intercept

def loss_fn(params, x, y):
    """MSE loss for one city."""
    predictions = predict(params, x)
    return jnp.mean((predictions - y) ** 2)

def train_step(params, x, y, learning_rate):
    """One gradient descent step for one city."""
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    new_params = params - learning_rate * grads
    return new_params, loss

Step 3: Vectorize Across Cities

Now we use vmap to run training for all three cities in parallel:

# Vectorize the training step
# params: axis 0 (each city has its own params)
# x: axis 0 (each city has its own data)
# y: axis 0 (each city has its own targets)
# learning_rate: None (same for all)
parallel_train_step = jax.jit(
    jax.vmap(train_step, in_axes=(0, 0, 0, None))
)

# Initialize parameters for all 3 cities
# Shape: (3, 2) - 3 cities, 2 params each (slope, intercept)
key, subkey = random.split(key)
params = random.normal(subkey, (n_cities, 2))

print("Initial parameters:")
for i in range(n_cities):
    print(f"  City {i}: slope={params[i, 0]:.2f}, intercept={params[i, 1]:.2f}")

Step 4: Train All Models in Parallel

learning_rate = 0.0001
n_epochs = 2000

print("\nTraining...")
for epoch in range(n_epochs):
    # This single line trains ALL THREE models simultaneously
    params, losses = parallel_train_step(params, X, Y, learning_rate)

    if epoch % 500 == 0:
        print(f"Epoch {epoch:4d} | Losses: {losses}")

print("\nFinal learned parameters:")
print(f"{'City':<6} {'True Slope':<12} {'Learned Slope':<15} {'True Intercept':<15} {'Learned Intercept':<15}")
print("-" * 70)
for i in range(n_cities):
    print(f"{i:<6} {true_slopes[i]:<12.1f} {params[i, 0]:<15.2f} {true_intercepts[i]:<15.1f} {params[i, 1]:<15.2f}")

Expected output:

Final learned parameters:
City   True Slope   Learned Slope   True Intercept   Learned Intercept  
----------------------------------------------------------------------
0      50.0         49.87           100.0            102.34         
1      30.0         29.92           200.0            201.15         
2      80.0         79.78           50.0             52.89

We trained three separate models, and there's not a single Python for-loop in the training logic. The parallel_train_step function processes all cities in one fused GPU kernel.

Why This Matters

The pattern we just used scales to serious applications:

Hyperparameter search: Train 100 models with different learning rates simultaneously. Pick the best one.

Ensemble methods: Train 10 models with different random seeds. Average their predictions for more robust results.

Per-user personalization: Train a tiny model for each of your 10,000 users. vmap handles the parallelization.

Bayesian methods: Sample 1000 parameter configurations from a posterior distribution and evaluate all of them at once.

The key insight is that vmap doesn't just save you from writing loops; it enables computations that would be impractical with sequential processing.

Exercises

1. Gradient verification: Use jax.grad to compute the derivative of f(x) = sin(x). Plot it alongside cos(x) to verify they match.

2. The ensemble challenge: Modify the project to train 10 models on the same city but with different random initializations. Use vmap over the params axis only (in_axes=(0, None, None, None)). Check if they all converge to similar values.

3. Second derivatives: jax.grad returns a function; and you can take its gradient too. Compute the second derivative of f(x) = x³ and verify it equals 6x.

Quick Reference

import jax
import jax.numpy as jnp

# vmap: Automatic Vectorization
def single_fn(x):
    return x ** 2

batched_fn = jax.vmap(single_fn)
results = batched_fn(jnp.array([1, 2, 3]))

# With multiple arguments
def dot(w, x):
    return jnp.dot(w, x)

# Shared weights, batched inputs
batch_dot = jax.vmap(dot, in_axes=(None, 0))

# grad: Automatic Differentiation
def loss(params):
    return params ** 2

grad_fn = jax.grad(loss)
gradient = grad_fn(3.0)  # 6.0

# Get both value and gradient
loss_val, grad_val = jax.value_and_grad(loss)(3.0)

# Gradient with respect to specific argument
def f(x, y):
    return x * y

df_dy = jax.grad(f, argnums=1)

# Combining Transforms
fast_batched_grad = jax.jit(jax.vmap(jax.grad(loss_fn), in_axes=(None, 0)))

What's Next

We've now covered the core JAX transforms: jit for speed, vmap for batching, and grad for gradients. These three tools are enough to train neural networks from scratch.

But writing raw JAX for complex models gets tedious. Next week, we'll introduce Flax NNX'; a neural network library that gives you PyTorch-style ergonomics while keeping all the power of JAX transformations.

We'll build our first real neural network: a CNN for image classification.

Resources

Automatic Vectorization

JAX AI Stack Guide

import numpy as np

x = np.random.randn(1000, 1000)
y = np.random.randn(1000, 1000)
result = np.dot(x, y)

NumPy is comfortable. It's the first thing we reach for when we need to do math. But here's the uncomfortable truth: for the kind of work we're about to do: training neural networks, processing millions of samples, running on GPUs and TPUs, NumPy is holding us back.

Not because NumPy is bad. It's genuinely excellent at what it was designed for. The problem is that NumPy was designed in 2005, before GPUs became the workhorses of machine learning, before TPUs existed, and before we needed to compute gradients of functions with millions of parameters.

JAX is what NumPy would look like if we designed it today, knowing what we know now.

By the end of this article, we'll have:

1. Understood why JAX is faster (not just that it is)

2. Written our first JAX code and seen the speedup ourselves

3. Learned the one mental shift that trips up everyone coming from NumPy or PyTorch

4. Built a working benchmark that proves the difference

Let's get into it.

The Problem with "Normal" Python

When you run Python code, the interpreter reads your instructions one line at a time, translates each one to machine code, executes it, then moves to the next line. This is called interpretation.

Line 1 → Translate → Execute
Line 2 → Translate → Execute
Line 3 → Translate → Execute
...

For a script that processes a CSV file or serves a web page, this is fine. The overhead of interpretation is negligible compared to the actual work being done.

But matrix multiplication? That's different. When we multiply two 5000×5000 matrices, we're doing 125 billion floating-point operations. The "translate → execute" overhead for each operation adds up fast.

NumPy helps by pushing the heavy lifting into compiled C code. When you call np.dot(), Python hands off the work to optimized BLAS libraries that run at near-hardware speed. That's why NumPy is faster than pure Python.

But there's still a problem: Python is still orchestrating the operations. Every time you chain NumPy calls together: np.dot(), then np.sum(), then np.exp(), Python has to:

1. Call into C

2. Wait for the result

3. Copy the result back to Python

4. Call into C again for the next operation

Each of those handoffs has overhead. And when you're doing this millions of times in a training loop, it adds up.

How JAX Fixes This: XLA Compilation

JAX takes a different approach. Instead of executing operations one at a time, JAX can compile your entire function into a single optimized program using XLA (Accelerated Linear Algebra).

Here's what that means in practice:

Read entire function → Analyze → Optimize → Fuse operations → Execute once

When JAX compiles a function, it:

• Fuses operations: Instead of computing a + b, storing the result, then computing result * c, XLA fuses these into a single kernel that does (a + b) * c without intermediate storage.

• Eliminates dead code: If you compute something but never use it, XLA removes it entirely.

• Optimizes memory access: XLA reorders operations to minimize cache misses and memory transfers.

• Targets your hardware: The same JAX code compiles to optimized instructions for CPU, GPU, or TPU without you changing anything.

The result? Code that runs 10x, 100x, sometimes 1000x faster than the NumPy equivalent.

Setting Up

Let's stop talking and start coding. We'll use Google Colab for this, it's free and gives us access to GPUs.

# Install the JAX AI stack (includes JAX, Flax, Optax, and friends)
!pip install -q jax-ai-stack

Now let's verify our setup:

import jax
import jax.numpy as jnp
import numpy as np
import time

print(f"JAX version: {jax.__version__}")
print(f"Available devices: {jax.devices()}")

If you're on a GPU runtime, you should see something like:

JAX version: 0.8.0
Available devices: [CudaDevice(id=0)]

If you see CpuDevice, that's fine too; JAX still provides speedups on CPU through XLA compilation.

The First Mental Shift: Immutability

Before we benchmark anything, we need to talk about the one thing that trips up everyone coming from NumPy or PyTorch.

In JAX, arrays are immutable. You cannot modify them in place.

In NumPy, this is perfectly normal:

# NumPy: Mutable arrays
arr = np.zeros(5)
arr[0] = 42  # Modify in place
print(arr)   # [42. 0. 0. 0. 0.]

In JAX, this will raise an error:

# JAX: This will FAIL
arr = jnp.zeros(5)
arr[0] = 42  #TypeError: JAX arrays are immutable

Why? Because immutability is what makes JAX's optimizations possible. If arrays can be modified from anywhere in your code, the compiler can't safely reorder operations or run them in parallel. By guaranteeing that arrays never change, JAX can aggressively optimize your code.

So how do we update arrays? We use the .at[].set() syntax, which returns a new array with the modification:

# JAX: The correct way
arr = jnp.zeros(5)
new_arr = arr.at[0].set(42)

print(arr)      # [0. 0. 0. 0. 0.] — Original unchanged
print(new_arr)  # [42. 0. 0. 0. 0.] — New array with the update

This feels wasteful at first, are we really copying the entire array just to change one element? In practice, no. JAX and XLA are smart enough to optimize this. But conceptually, you should think of it as creating a new array.

Here's the full set of .at[] operations:

x = jnp.array([1, 2, 3, 4, 5])

# Set a value
x.at[0].set(10)         # [10, 2, 3, 4, 5]

# Add to a value
x.at[0].add(10)         # [11, 2, 3, 4, 5]

# Multiply a value
x.at[0].multiply(10)    # [10, 2, 3, 4, 5]

# Works with slices too
x.at[1:3].set(99)       # [1, 99, 99, 4, 5]

Commit this to memory. You'll use it constantly.

The Benchmark: NumPy vs JAX vs JAX+JIT

Now let's prove that JAX is actually faster. We'll multiply two large matrices; this is the core operation in neural networks (every linear layer is a matrix multiplication).

Step 1: Create the Data

# Matrix size
size = 3000

# Create random matrices with NumPy
x_np = np.random.normal(size=(size, size)).astype(np.float32)
y_np = np.random.normal(size=(size, size)).astype(np.float32)

# Convert to JAX arrays
x_jax = jnp.array(x_np)
y_jax = jnp.array(y_np)

print(f"Matrix shape: {x_np.shape}")
print(f"Total elements per matrix: {size * size:,}")
print(f"Operations for multiplication: {size ** 3:,}")

Output:

Matrix shape: (3000, 3000)
Total elements per matrix: 9,000,000
Operations for multiplication: 27,000,000,000

That's 27 billion operations. Let's see who can do it fastest.

Step 2: Benchmark NumPy

def matmul_numpy(x, y):
    return np.dot(x, y)

# Warmup
_ = matmul_numpy(x_np, y_np)

# Timed run
start = time.perf_counter()
result_np = matmul_numpy(x_np, y_np)
numpy_time = time.perf_counter() - start

print(f"NumPy time: {numpy_time:.4f} seconds")

Result:

NumPy time: 0.6294 seconds

Step 3: Benchmark JAX (without JIT)

def matmul_jax(x, y):
    return jnp.dot(x, y)

# Warmup
_ = matmul_jax(x_jax, y_jax).block_until_ready()

# Timed run
start = time.perf_counter()
result_jax = matmul_jax(x_jax, y_jax).block_until_ready()
jax_time = time.perf_counter() - start

print(f"JAX time (no JIT): {jax_time:.4f} seconds")

Important: We call .block_until_ready() because JAX operations are asynchronous. When you call jnp.dot(), JAX immediately returns a "future" and continues executing Python code while the GPU works in the background. Without block_until_ready(), we'd be timing how fast JAX can dispatch the operation, not how fast it actually runs.

Result:

JAX time (no JIT): 0.0094 seconds

Step 4: Benchmark JAX with JIT Compilation

Here's where JAX shows its true power. We add a single decorator:

@jax.jit
def matmul_jax_jit(x, y):
    return jnp.dot(x, y)

# First call: JAX traces and compiles the function
# This takes a moment, so we don't include it in the benchmark
print("Compiling...")
_ = matmul_jax_jit(x_jax, y_jax).block_until_ready()
print("Done.")

# Timed run (using the compiled version)
start = time.perf_counter()
result_jit = matmul_jax_jit(x_jax, y_jax).block_until_ready()
jit_time = time.perf_counter() - start

print(f"JAX time (with JIT): {jit_time:.4f} seconds")

The first call to a @jax.jit function is slow because JAX is tracing your function, figuring out what operations it contains, and then compiling it with XLA. Subsequent calls use the compiled version and are extremely fast.

Step 5: Compare Results

print("RESULTS")
print(f"NumPy:          {numpy_time:.4f} seconds")
print(f"JAX (no JIT):   {jax_time:.4f} seconds")
print(f"JAX (with JIT): {jit_time:.4f} seconds")
print(f"\nSpeedup (JIT vs NumPy): {numpy_time / jit_time:.1f}x")

Typical results on a Colab GPU:

RESULTS
NumPy:          0.6294 seconds
JAX (no JIT):   0.0094 seconds
JAX (with JIT): 0.0198 seconds

Speedup (JAX vs NumPy):     67.2x
Speedup (JIT vs NumPy):     31.8x
Speedup (JIT vs JAX):       0.5x

What Just Happened?

Let's break down why the JIT version is so much faster:

1. Hardware acceleration: JAX moved the computation to the GPU, which has thousands of cores optimized for parallel math.

2. XLA compilation: Even on GPU, the JIT version is faster than raw JAX because XLA fuses operations and optimizes memory access patterns.

3. No Python overhead: Once compiled, the function runs entirely in native code. Python is only involved in dispatching the call.

The key insight is that @jax.jit doesn't just run your code on a GPU; it fundamentally changes how your code runs.

The Randomness Trap (Bonus Lesson)

There's one more gotcha that catches everyone early on. Try this:

@jax.jit
def broken_random():
    return np.random.randn(5)  # Using NumPy's random

result1 = broken_random()
result2 = broken_random()
print(f"First call:  {result1}")
print(f"Second call: {result2}")

You'll notice that result1 and result2 are identical. The random numbers got "baked in" during compilation.

JAX requires explicit random state management. Here's the correct way:

from jax import random

@jax.jit
def correct_random(key):
    return random.normal(key, shape=(5,))

# Create a PRNG key
key = random.PRNGKey(42)

# Split the key for each use
key, subkey1 = random.split(key)
result1 = correct_random(subkey1)

key, subkey2 = random.split(key)
result2 = correct_random(subkey2)

print(f"First call:  {result1}")
print(f"Second call: {result2}")

Now you get different random numbers each time. We'll cover this pattern in depth when we build neural networks, but for now, just remember: never use np.random inside JIT-compiled functions.

Exercises

Before moving on, try these:

1. Break the rules: Try to modify a JAX array in place (x[0] = 1). Read the error message carefully: JAX errors are verbose but informative.

2. Vary the size: Run the benchmark with different matrix sizes (1000, 2000, 5000). How does the speedup change?

3. Chain operations: Write a function that does multiple operations (jnp.dot, then jnp.sum, then jnp.exp). Compare JIT vs non-JIT. The speedup should be even larger because XLA fuses the operations.

What's Next

We've established why JAX is fast and seen the proof. But speed is only half the story. Next week, we'll explore transformations; the features that make JAX genuinely different from NumPy, not just faster.

Specifically, we'll cover:

• jax.vmap: Automatic vectorization that eliminates for-loops

• jax.grad: Automatic differentiation that makes backpropagation trivial

These two functions are why JAX has become the framework of choice for machine learning research. Once you understand them, you'll never look at NumPy the same way again.

Quick Reference

import jax
import jax.numpy as jnp
from jax import random

# Basic array operations (same as NumPy)
x = jnp.array([1, 2, 3])
y = jnp.zeros((3, 3))
z = jnp.dot(a, b)

# Updating arrays (immutable style)
new_x = x.at[0].set(99)
new_x = x.at[1:].add(10)

# JIT compilation
@jax.jit
def fast_function(x):
    return jnp.dot(x, x.T)

# Explicit randomness
key = random.PRNGKey(0)
key, subkey = random.split(key)
samples = random.normal(subkey, shape=(100,))

# Block for accurate timing
result = fast_function(x).block_until_ready()

Next week: Transformations That Change Everything—Automatic Vectorization with vmap and Gradients with grad

Resources:

JAX Documentation

JAX AI Stack

Notebook

It sounds absurd, but the "Reverse Turing Test" is here. Writers are deliberately inserting typos, breaking grammatical rules, and shunning perfectly functional words like "delighted," "landscape," or "delve," replacing them with awkward alternatives. Why? Because the black-box algorithms of AI detection tools have decided that high-proficiency English is evidence of a robot.

We have reached a dangerous inflection point where linguistic competence is treated with suspicion, and human error is fetishized as the only remaining proof of authenticity. But should the degradation of language really be the price of being believed?

The statistical mirror: Why AI sounds like us

To understand this crisis, we must first demystify the adversary. Large Language Models (LLMs) like GPT-4 or Gemini are not sentient poets; they are probabilistic engines. They are trained on the internet’s vast corpus of text; trillions of words written by humans over centuries.

When an AI uses words like "pivotal," "crucial," or "emphasize," it is not because it has a personal preference for corporate speak. It is demonstrating Zipf’s Law. This linguistic principle states that in any natural language, a small number of words are used with disproportionately high frequency to maximize efficiency. Humans naturally gravitate toward words that reduce cognitive load while maintaining clarity. LLMs simply mirror this statistical reality.

Therefore, the list of "banned" AI-sounding words currently circulating on social media reads like the standard vocabulary of any Ugandan NGO report, government white paper, or academic thesis from the last twenty years. If AI sounds "academic," it is only because it was trained on our academia. To penalize a writer for using structure and precision is to penalize them for being well-read.

The bias against excellence

The reliance on AI detectors is not just unscientific; it is discriminatory. We are witnessing a collision with Goodhart’s Law: "When a measure becomes a target, it ceases to be a good measure." By using "perplexity" (a measure of randomness) to judge humanity, detectors punish clear, logical writing.

This has grave implications for Africa. A 2023 study by researchers at Stanford University revealed a stunning bias: AI detectors flagged over 61% of essays written by non-native English speakers as AI-generated, compared to nearly zero for native US 8th graders.

For the African student who has spent years mastering the "Queen’s English"; learning the formal transitions and structured arguments prized by schools, this is a slap in the face. Writing with the clarity and structure taught in our schools now puts you at risk of being labeled a fraud. The message sent to our students and professionals is chilling: Write badly, or be doubted.

The soul in the machine

However, while we should not dumb down our syntax, we must accept that AI forces us to elevate our substance.

AI can mimic the form of human expression, the rhythm of a sonnet or the structure of a press release, but it lacks the referent. It has no connection to the physical world. It processes symbols, not reality.

This is where the true differentiation lies. An AI can generate a paragraph about the concept of "cultural heritage," but it does not know the specific, heavy silence that falls over a clan meeting when obuntu bulamu is violated. It can describe the ingredients of luwombo, but it cannot understand the politics of the banana plantation or why the preparation of food is a language of love in Buganda.

AI operates on prediction; humans operate on intention.

Contextual wisdom: A model knows that traffic jams are bad. It does not know the specific, communal frustration of a Friday evening gridlock on Jinja road, nor the humor exchanged between strangers in a taxi.

Emotional weight: AI can string together words about grief, but it cannot choose a metaphor that breaks the heart because it has never had a heart to break. It cannot anchor a sentence in the lived experience of paying school fees in January.

The path forward

We must reject the impulse to perform incompetence. Trying to "beat the detector" by inserting errors is a losing battle; the models will eventually learn those tricks too.

Instead, the rise of AI should force a renaissance of voice. The era of generic, filler-heavy writing is indeed over; not because it is "AI," but because AI can do it faster. The human writer must now bring something the machine cannot: original insight grounded in lived reality.

We must lean into our idiosyncrasies, our cultural nuances, our irony, and our specific Ugandan perspectives. We must tell stories that rely on the messy, unpredictable texture of real life.

The real test of humanity is not whether we can write "less like a robot." It is whether we can think, feel, and observe the world deeply enough to say something that no statistical model could ever predict. On that front, we still hold the advantage.

Responding to years of passionate community requests manifested in blog posts, forum threads, and creative GitHub workarounds, Google has officially bridged this gap. Today, we are thrilled to explore the new Google Colab extension for Visual Studio Code, a tool that promises the best of both worlds. This article provides a detailed, technical tutorial on how to install, configure, and harness the power of this game-changing extension, transforming your local VS Code into a control room for Google’s heavy-lifting cloud infrastructure.

The Best of Both Worlds: Unifying Local IDE and Cloud Compute

The core value of the Colab extension is its ability to meet developers where they are. It acknowledges that while Colab’s simplicity is a major strength, many users crave the advanced features of a full-fledged IDE for larger projects and complex workflows.

• For VS Code Users: The primary advantage is the ability to connect local .ipynb notebooks to high-powered Colab runtimes. This means you can continue using your familiar, highly customized editor while seamlessly accessing premium GPUs and TPUs, including those available through Colab Pro subscriptions, without leaving your local environment.

• For Colab Users: This integration supports the common practice of working on notebooks that are part of a larger project or Git repository. It empowers users who need more robust IDE features—such as superior code completion, version control, and advanced debugging—by pairing the simplicity of Colab's provisioned runtimes with the prolific VS Code editor.

Essentially, this move bridges the gap between code productivity and cloud compute scalability, eliminating the need to switch tabs, export notebooks, or manage credentials across different platforms.

Getting Started: A Step-by-Step Guide

You can get up and running with the Colab extension in just a few clicks. The setup is designed to be intuitive and fast.

Step 1: Install the Colab Extension

First, you need to add the extension to your VS Code installation.

1. Open the Extensions view from the Activity Bar on the left side of your VS Code window (or press Ctrl+Shift+X).

2. In the marketplace search bar, type Google Colab.

3. Click Install on the official extension published by Google.

4. If you do not already have it, the installer will prompt you to install its required dependency, the official Jupyter extension.

Step 2: Connect to a Colab Runtime

Once installed, you can connect any local notebook to a Colab runtime.

1. Create a new notebook (.ipynb file) or open an existing one in your local VS Code workspace.

2. To select the execution environment, you can either run a cell, which will prompt you to choose a kernel, or click the Select Kernel button in the top-right corner of the notebook interface.

3. From the dropdown menu, choose Select Another Kernel...

4. Click on the Colab option. You will be prompted to sign in with your Google account.

5. After signing in, you can choose to create a New Colab Server or connect to an existing one you may have running. For your first time, you will create a new one.

Your local notebook is now powered by a Google Colab runtime! You can give your Colab server a name so that it can be easily referenced and reused in the future.

Step 3: Select Your Compute Resources

The true power of this extension lies in accessing specialized hardware. The available accelerator options and memory limits are determined by your Google Colab subscription plan.

• Free Tier: Users have access to NVIDIA T4 GPUs and TPU v5e accelerators.

• Colab Pro Tier: Subscribers gain access to more powerful hardware, such as premium GPUs like the NVIDIA A100.

After connecting to the Colab kernel, you can select your desired hardware accelerator for the session.

Practical Examples — Beyond the Basics

To truly appreciate this new workflow, let's move beyond simple demonstrations. Here are two original, real-world scenarios that are impractical on a standard local machine but become trivial with the Colab extension.

Example 1: GPU-Accelerated Big Data Analysis with RAPIDS cuDF

The Challenge: You need to analyze a large CSV file (several gigabytes) containing millions of records. Using a standard library like Pandas on a CPU can be painfully slow, with simple grouping and aggregation operations taking minutes to complete.

The Solution: We'll use RAPIDS cuDF, a GPU-accelerated DataFrame library with a Pandas-like API. By running this in VS Code connected to a Colab GPU, we can perform the analysis in seconds. RAPIDS cuDF is now pre-installed in Colab GPU runtimes, making this seamless.

The Result:

The complex aggregation on 3,066,766 rows of data completes in just 0.3070 seconds. This incredible speed, demonstrated in the output below, transforms what would be a coffee-break task on a CPU into an interactive, real-time query. This showcases a real-world data engineering task made efficient and seamless, all within the comfort of VS Code.

The Code:

You can find a copy of the corresponding Colab workbook for this example here.

Example 2: Creative AI — Generating Art with Stable Diffusion

The Challenge: Text-to-image models like Stable Diffusion are computationally expensive and require significant GPU VRAM, making them inaccessible to users without high-end local hardware.

The Solution: We'll use the Hugging Face diffusers library to run a Stable Diffusion pipeline on our Colab GPU kernel. This allows us to generate high-quality images from text prompts directly inside a VS Code notebook.

Install required libraries We need diffusers, transformers, and accelerate for this task.

Set up the Stable Diffusion Pipeline This code downloads the pre-trained model weights and prepares the pipeline for inference on the GPU.

Generate an Image Define your creative prompt and let the model generate an image.

The Result: Within a minute, a high-resolution, AI-generated image appears directly in your VS Code notebook output. This showcases how the extension democratizes access to powerful generative AI models, enabling creative experimentation without the need for a dedicated local GPU.

The Code:

You can find the corresponding Colab notebook for this example here.

Advanced Tips and Current Limitations

To get the most out of the extension, keep these points in mind:

• File Management: You are working on a remote Colab file system. Use commands like !ls -l or VS Code's built-in file explorer to see files generated in your runtime. To persist your work, consider mounting your Google Drive:

    from google.colab import drive
    drive.mount('/content/drive')
  

• Secrets Management: The web UI's native secrets manager is not yet available. For securely handling API keys, use the file-upload workaround to upload a .env file, as detailed in the source articles.

• Session Lifetime: Remember that Colab runtimes are ephemeral. They will disconnect after a period of inactivity (typically 90 minutes for free-tier users) or if you exceed the maximum session duration (12 hours). Save your work frequently.

The Bigger Picture: What's Next?

As a newly released tool, the Colab extension is still in its early stages, and some limitations exist. As noted, certain web-UI-specific functions like the secrets manager are not yet implemented. However, Google has positioned this release as a "launchpad," signaling a commitment to bringing even more of Colab's functionality to developers everywhere.

This launch also places Google in a fascinating strategic position, turning VS Code into a key battleground for AI developer mindshare. By bringing its powerful code execution engine into the same interface where tools like GitHub Copilot excel at code generation, Google is challenging the AI-assisted developer landscape.

For developers, this rising competition is a net positive. It promises a future where the lines between code generation and execution blur, and where powerful, integrated, and accessible AI tools become a fundamental component of the editor itself.

Conclusion

The new Google Colab extension for VS Code is more than just a convenience; it's a transformative tool that unifies the best of local development with the power of cloud computing. It empowers developers and ML engineers to harness free GPUs and TPUs directly within their preferred editor, streamlining workflows and accelerating innovation. While still in its early stages, the extension represents a significant step forward in making AI and machine learning development more accessible and productive. The future looks bright, and it runs on a seamless connection between your local machine and the cloud.

Our agent will demonstrate how to:

• Combine LLMs and Tools: Overcome LLM limitations by letting them call Python functions (tools) for fresh data.

• Build Custom Tools: Write scrapers and generators (e.g. for attractions and images) that the agent can invoke.

• Use Gemini as the LLM: Leverage Google’s multimodal Gemini 2.5 Pro model (our most advanced AI model) for reasoning.

• Leverage LlamaIndex: Use LlamaIndex’s ReActAgent and FunctionTool wrappers to glue tools and the LLM together.

• See the ReAct Loop in Action: Peek at the agent’s chain of thought as it solves a query by sequentially “thinking” and “acting”.

• Advanced Prompting: Go beyond raw data by feeding tool outputs into a carefully crafted prompt for insightful recommendations.

By the end, you’ll have a working Tourism AI Assistant that answers travel questions, scrapes live data, finds images, and even tweets about destinations—all powered by Gemini and LlamaIndex.

Our Technology Stack

Before coding, let’s understand the pieces we’ll use:

• Google Gemini 2.5 Pro: A state-of-the-art multimodal LLM (text+code+images+video) with strong reasoning capabilities. This is the “brain” of our agent. We use the models/gemini-2.5-pro endpoint via Google’s AI Generative SDK.

• LlamaIndex: An open-source framework for building LLM apps. It provides the ReActAgent class and FunctionTool wrapper to connect LLMs with Python tools. LlamaIndex lets us expose any function to the agent in a structured way.

• Web Scraping Libraries: We’ll use requests and beautifulsoup4 to fetch live data from travel sites (for attractions, etc.). (In production, a tourism API is often preferred to avoid brittle scrapers, but for our tutorial we’ll show how to scrape responsibly.)

• Image Search: We’ll implement a simple DuckDuckGo image scraper to let the agent fetch pictures of destinations.

• Content Generator: A function to draft a promotional social-media post about a place or event, illustrating how the agent can take action, not just fetch data.

Installing Dependencies

First, install the required Python packages. In your notebook or script, run:

!pip install -q llama-index llama-index-llms-gemini google-generativeai python-dotenv beautifulsoup4 requests

• llama-index-llms-gemini: Adds Gemini support to LlamaIndex.

• google-generativeai: Google’s SDK to call the Gemini API.

• python-dotenv: For loading API keys from a .env file (or Colab secrets).

• beautifulsoup4 & requests: For our web-scraping tools.

We’ll also enable detailed logging to trace the agent’s reasoning:

import os, logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

Authentication

To use Gemini, you need a Google AI Studio API key:

1. Create an API key in Google AI Studio.

2. Store it as GOOGLE_API_KEY in your environment (or Colab secrets).

Then configure the Google Generative AI SDK:

import google.generativeai as genai
from dotenv import load_dotenv

load_dotenv()  # if using .env file
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
if not GOOGLE_API_KEY:
    raise ValueError("Please set GOOGLE_API_KEY environment variable.")
genai.configure(api_key=GOOGLE_API_KEY)
logger.info("Configured Google API key.")

Configuring the LLM and Embedding Model

We tell LlamaIndex which LLM (Gemini) and embedding model to use. The embedding model isn’t crucial here (we won’t do retrieval-heavy tasks), but LlamaIndex requires one. We’ll use a small HuggingFace model:

from llama_index.llms.gemini import Gemini
from llama_index.core import Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

MODEL_NAME = "models/gemini-2.5-pro"  # Google’s Gemini 2.5 Pro model

logger.info("Configuring LlamaIndex settings...")
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.llm = Gemini(model=MODEL_NAME)  # uses GOOGLE_API_KEY automatically
logger.info(f"Using Gemini LLM: {MODEL_NAME}")

Note: Gemini 2.5 Pro is a powerful model that excels at complex tasks. It supports text, code, and images, and even has built-in “thinking” ability for chain-of-thought reasoning.

Crafting Our Tools

An agent is only as capable as its tools. A tool here is a Python function that performs some action—like fetching data or creating content—that the LLM alone cannot do (e.g. real-time web queries). We’ll create a small toolkit for tourism tasks:

• Tool 1: get_city_attractions(city) – Scrape top attractions for a city. (As an example, we’ll scrape a known travel guide or Wikipedia for a few highlights.)

• Tool 2: get_city_restaurants(city) – (Optional) Scrape or list popular restaurants. For brevity, this could return a static list or use a simple scrape.

• Tool 3: search_for_destination_images(place) – Search DuckDuckGo for images of a place (like “Eiffel Tower Paris”). Returns a few image URLs.

• Tool 4: generate_tourism_post(place, summary) – Generate a friendly promotional social-media post (tweet) about an attraction or city, given its name and a short summary.

Each tool will have a clear docstring explaining its purpose and inputs. The agent reads these docstrings to know when to use which tool. For example:

import requests
from bs4 import BeautifulSoup

def get_city_attractions(city: str) -> str:
    """
    Fetches the top tourist attractions for the given city.
    Returns a bullet list of attractions and brief info.
    """
    try:
        # Example: scrape PlanetWare or Wikipedia page for top attractions
        url = f"https://www.planetware.com/{city.lower()}/top-rated-tourist-attractions-in-{city.lower()}.htm"
        resp = requests.get(url)
        soup = BeautifulSoup(resp.text, "html.parser")
        attractions = [h2.get_text(strip=True) for h2 in soup.find_all('h2')[:5]]
        if not attractions:
            raise Exception("No attractions found")
        return "\n".join(f"- {a}" for a in attractions)
    except Exception as e:
        return f"Could not retrieve attractions for {city} (error: {e})"

def get_city_restaurants(city: str) -> str:
    """
    Returns a short list of popular restaurants in the city.
    (For demo purposes, this may be a static list or scraped from a simple source.)
    """
    # Placeholder: in a real app, use an API or reliable source
    dummy_data = {
        "Paris": ["Le Jules Verne (Eiffel Tower)", "L'Ambroisie", "Septime"],
        "Rome": ["Da Enzo al 29", "Roscioli", "La Pergola"]
    }
    return "\n".join(f"- {r}" for r in dummy_data.get(city, ["No data available"]))

def search_for_destination_images(query: str) -> list:
    """
    Searches DuckDuckGo Images for the query and returns a list of image URLs.
    """
    try:
        res = requests.get(f"https://duckduckgo.com/?q={query.replace(' ', '+')}&iar=images&iax=images")
        soup = BeautifulSoup(res.text, "html.parser")
        imgs = soup.select("img.tile--img__img")[:5]
        return [img.get("src") for img in imgs if img.get("src")]
    except Exception as e:
        return [f"Image search failed: {e}"]

def generate_tourism_post(place: str, summary: str) -> str:
    """
    Generates a friendly social media post about a place.
    """
    prompt = f"Write an enthusiastic tweet about visiting {place}. Summary: {summary}"
    resp = genai.generate_message(model=MODEL_NAME, messages=[{"role": "user", "content": prompt}])
    return resp.last

Note: In production, scraping arbitrary sites (like PlanetWare) can be fragile. Sites may block scrapers, change layout, or have legal terms against scraping. For robust applications, a dedicated travel data API is preferable. Here we use simple scrapers for illustration.

Wrapping Functions for LlamaIndex

LlamaIndex needs tools wrapped in FunctionTool objects. This exposes each function’s signature and docstring to the LLM. Docstrings become part of the agent’s understanding of when to use each tool. For example:

from llama_index.core.tools import FunctionTool

logger.info("Wrapping functions into LlamaIndex tools...")
tools = [
    FunctionTool.from_defaults(fn=get_city_attractions,
                               name="get_city_attractions",
                               description="Get the top tourist attractions in a city."),
    FunctionTool.from_defaults(fn=get_city_restaurants,
                               name="get_city_restaurants",
                               description="Get popular restaurants in a city."),
    FunctionTool.from_defaults(fn=search_for_destination_images,
                               name="search_for_destination_images",
                               description="Search for images of a place."),
    FunctionTool.from_defaults(fn=generate_tourism_post,
                               name="generate_tourism_post",
                               description="Create a social media post promoting an attraction or city.")
]
logger.info(f"Created {len(tools)} tools: {[t.name for t in tools]}")

Each FunctionTool includes the function name, parameters, and a human-readable docstring. The agent will choose which tool to call by looking at your query and the tool descriptions, then pass the right arguments.

Initializing the ReAct Agent

With tools ready, we create the agent. We use LlamaIndex’s ReActAgent, which implements the Reason-and-Act loop. The agent will think about the user’s question, decide which tool to use and with what inputs, execute the tool, then observe the result and repeat as needed. Setting verbose=True lets us see the entire chain of thought:

from llama_index.core.agent import ReActAgent

logger.info("Initializing the ReAct agent...")
agent = ReActAgent.from_tools(tools=tools, llm=Settings.llm, verbose=True)
logger.info("Tourism Agent is ready to go!")

Under the hood, ReActAgent will prompt Gemini with a system message describing the tools and this reasoning framework. The LLM will follow a structure like:

Thought: I should use [tool] because ...
Action: [tool] with input {...}
Observation: (tool output)
Thought: ... continue or final answer.

This explicit reasoning flow is characteristic of the ReAct approach.

Trying Out the Tourism Agent

Our agent is now online! Let’s test it with a couple of example queries and inspect its reasoning (verbose output):

Scenario 1: Simple Question

User: “What are the top attractions in Paris?”

response = agent.chat("What are the top tourist attractions in Paris?")
print(response)

Agent’s (verbose) Thought Process:

• Thought: The user asks for top attractions in Paris. The get_city_attractions tool seems appropriate.

• Action: get_city_attractions

• Action Input: {"city": "Paris"}

• Observation:

    - Eiffel Tower
    - Louvre Museum
    - Notre-Dame Cathedral
    - Sacré-Cœur Basilica
    - Musée d'Orsay
  

• Thought: I have the list of attractions. Now I answer the user using this info.

• Answer: “Top attractions in Paris include the Eiffel Tower, Louvre Museum, Notre-Dame Cathedral, Sacré-Cœur, and Musée d’Orsay.”

Agent Answer: Top attractions in Paris include: Eiffel Tower, Louvre Museum, Notre-Dame Cathedral, Sacré-Cœur Basilica, and Musée d’Orsay.

This shows the ReAct loop: the agent identified the need for a tool, executed it, and synthesized the output.

Scenario 2: Multi-Step Task

User: “Find art-related attractions in Paris and then write a tweet about them.”

This requires multiple steps: identify art attractions, then generate a tweet.

response = agent.chat("Find art-related attractions in Paris and draft a tweet about them.")
print(response)

Agent’s Thought Process:

• Thought: The user needs two things: 1) art-related attractions in Paris, 2) a tweet about them. First, get all attractions.

• Action: get_city_attractions

• Input: {"city": "Paris"}

• Observation: (same list as before)

• Thought: The attractions list is obtained. Now pick those related to art (e.g. museums, galleries). These include Louvre Museum and Musée d'Orsay. I will create a summary and use the generate_tourism_post tool.

• Action: generate_tourism_post

• Input: {"place": "Paris Museums", "summary": "Highlights: Louvre Museum and Musée d'Orsay"}

• Observation:

    Explore Paris's Artistic Treasures!
    Don’t miss the Louvre Museum’s timeless classics and the Musée d’Orsay’s Impressionist masterpieces. From ancient wonders to modern art, Paris has it all. Bon voyage and happy exploring!
  

• Thought: I have the tweet text. Now I can share the final answer.

• Answer: Here’s a tweet promoting Paris’s art attractions:… (tweet above)…

Agent Answer: “Here’s a tweet about Paris’s art attractions: Explore Paris’s Artistic Treasures! Don’t miss the Louvre Museum’s timeless classics and the Musée d’Orsay’s Impressionist masterpieces. From ancient wonders to modern art, Paris has it all. Bon voyage and happy exploring!”

This demonstrates chaining: the agent used get_city_attractions, filtered the results, then used generate_tourism_post with the summary. The LLM reasoned at each step, as expected in a ReAct agent.

Advanced Prompting: Personalized Recommendations

So far we’ve fetched raw data and posted about it. But LLMs can also provide insights on top of data. For example, suppose a user asks:

“I’m a history buff visiting Rome. Which attractions should I see?”

Simply returning a list of all Rome attractions isn’t ideal. Instead, we want personalized recommendations. We can achieve this by combining our tools with a clever prompt:

1. Get raw data via tool: Call get_city_attractions("Rome").

2. Engineer a focused prompt: Tell Gemini to read the list and pick the best 3-5 for a history enthusiast, explaining why.

3. Generate the answer with Gemini: The LLM acts as a reasoning layer on the data.

For example:

# Step 1: fetch attractions
attractions = get_city_attractions("Rome")
# (assume this returns a bullet list of sites like Colosseum, Vatican, Pantheon, etc.)

# Step 2: craft a targeted prompt
prompt = f"""
You are an expert tour guide for Rome. A user is interested in historical sites. 
From the following list of Rome attractions, recommend 3-5 must-see locations for a history buff. Explain why each is a good choice.
Rome Attractions:
{attractions}

Focus on historical significance and include brief descriptions.
"""
# Step 3: call Gemini directly with this prompt
historical_resp = genai.generate_message(model=MODEL_NAME, messages=[{"role": "user", "content": prompt}])
print(historical_resp.last)

Sample Output:

Excellent choices for a history enthusiast in Rome include:

• Colosseum – The iconic ancient amphitheater where gladiators once fought. A symbol of Rome’s imperial past and architecture.

• Roman Forum – The heart of ancient Rome’s political and social life. You can walk among ruins of temples and government buildings.

• Pantheon – A 2,000-year-old temple-turned-church, showcasing Rome’s engineering and dedication to the gods. Its monumental dome is a marvel.

• Vatican Museums & St. Peter’s Basilica – While a church, the Vatican holds vast historical and artistic treasures spanning millennia, including Raphael’s Rooms and Michelangelo’s Pietà.

• Catacombs of Callixtus – Underground burial chambers that reveal early Christian history and traditions.

Each of these sites offers rich historical insights into Rome’s past civilizations and will captivate any history buff!

This “reasoning layer” approach leverages Gemini’s understanding to analyze and filter the tool output. We transformed a raw list into a personalized recommendation list with explanation. One could even wrap this workflow into a new tool (e.g. recommend_attractions_for_interest) for the agent to use directly.

Conclusion

We’ve built a fully functional Tourism AI Assistant that can interpret user requests, choose the right tool, fetch live data, and present it in a helpful way. We saw:

• Setting up the agent: Configuring Gemini and LlamaIndex.

• Writing tools: Python functions with clear docstrings, wrapped as FunctionTool for the agent.

• Using ReAct: The agent’s chain of thought (“Thought”, “Action”, “Observation”) shows how it plans and executes.

• Multimodal capability: Although we didn’t show it here, Gemini supports images and code, so you could extend this agent to analyze photos of landmarks or compute routes.

• Advanced prompting: We enhanced raw data by feeding it into Gemini with a crafted prompt, yielding richer, customized advice.

This agentic architecture is highly extensible. Next steps might include integrating real travel APIs (for flights or hotels), adding a memory so the assistant recalls user preferences, or connecting to mapping services for directions. The key idea is that the LLM does the reasoning, while we supply specialized tools for any real-world data or action.

Happy travels and happy building!

So, who cares about the mental health of software engineers and IT professionals, particularly those powering Africa's digital transformation? The answer must be: we all should. Behind the celebratory headlines of Africa's tech unicorns, a human cost is mounting – the often-invisible erosion of mental well-being among the architects of this progress. This is not merely a collection of individual struggles; it is a systemic challenge that threatens the sustainability of innovation and the very fabric of our tech ecosystems.

As we observe Mental Health Awareness Month this May, it is time to move beyond acknowledging this crisis to implementing tangible changes. My experiences, and those of my colleagues across the continent, underscore the urgency of this call. The narrative of "Africa Rising" in technology is inspiring, but it must not overshadow the profound pressures faced by those on the ground. The expectation to contribute to this grand narrative, often with limited resources and against significant infrastructural odds, can itself become a significant stressor.

Beyond the Global North: The Unique Pressures on Africa's Tech Talent

Globally, the mental toll on tech professionals is well-documented. Demanding work environments, tight deadlines, the constant need for upskilling, and the specter of burnout, anxiety, and imposter syndrome are common parlance in Silicon Valley and beyond. Studies show that over half of software engineers globally report burnout. However, in Africa, these universal strains are not just replicated; they are amplified and compounded by a unique set of local realities.

In many Ugandan communities, and indeed across much of Africa, mental health remains a taboo subject. Seeking help for anxiety or depression can be misconstrued as personal weakness or even attributed to spiritual failings, rather than being recognized as a legitimate health concern. This cultural barrier is a formidable obstacle; if employees are afraid to speak out, companies may remain unaware of the depth of the problem or feel unmotivated to invest in solutions. This silence perpetuates a cycle where low demand for services (due to stigma) leads to low investment, reinforcing the notion that these are not critical workplace issues.

Economic and infrastructural realities add further pressure. In a context where a tech job can be a life-changing opportunity for an individual and their family, the stakes for success are incredibly high. The daily grind can be exacerbated by challenges like unreliable power or expensive and erratic internet connectivity – issues that can turn a straightforward coding task into a marathon of frustration and directly impact productivity and stress levels.

For those of us working in Machine Learning and AI in Africa, there are additional, nuanced burdens. The global discourse on "Ethical AI" and "AI for Good" often places a heavy responsibility on African engineers to solve complex societal problems, sometimes without adequate local data, resources, or established ethical frameworks. We grapple with building AI that is not only innovative but also culturally relevant and free from biases that could harm vulnerable populations, for example, in healthcare or financial services. The challenge of "decolonizing AI," ensuring our creations are truly beneficial and equitable for African contexts, adds a significant mental and ethical load. This isn't just about delivering a product; it's about navigating a complex ethical landscape with potentially far-reaching consequences, a unique form of pressure that can lead to "ethical burnout."

"System Error": The Current State of Mental Health Support in African Tech

When we look at the mental health support typically available within African tech companies and startups, it often resembles a patchwork of well-intentioned but ultimately inadequate measures. Employee Assistance Programs (EAPs), if they exist, may be generic and not culturally attuned. The approach is frequently reactive, addressing crises rather than fostering a culture of proactive well-being. Many "wellness" initiatives might touch on stress management but fail to address the deep-seated systemic issues or the unique pressures of the tech environment.

Consider a software developer in Kampala, let's call her Aisha. Aisha is brilliant and dedicated, working for a promising fintech startup. She's battling intense burnout from months of 70-hour weeks leading up to a critical product launch. When she finally musters the courage to hint at her struggles to her manager, she’s met with a well-meaning but ultimately unhelpful suggestion to "take a short break" and a reminder of how crucial her role is. There's no formal HR pathway for mental health support, no access to confidential counselling that understands her context. Her story, and others like it – the junior IT support staff in Nairobi feeling overwhelmed by user demands and job insecurity, or the data scientist in Lagos wrestling with the ethical implications of an algorithm with little institutional guidance – are common, though often whispered in hushed tones. These anonymized experiences highlight a gap: the tech sector, while creating innovative solutions for Africa, sometimes fails to apply that same innovative spirit to supporting its own people.

There are glimmers of hope. Organizations like Mindverse Uganda are conducting workshops in companies, and innovative approaches like the Tele-Support Psychotherapy (TSP) model show the potential of culturally sensitive, tech-enabled solutions, even if primarily focused outside the corporate sphere for now. Some forward-thinking companies, like Andela or Interswitch, have begun to list mental health support among their employee perks. However, these are often exceptions. The "startup culture" globally tends to prioritize breakneck growth over sustained well-being, a trend likely amplified in Africa's resource-tighter, higher-pressure emerging ecosystems. This can lead to a systemic lack of support, masked by a few positive but isolated examples.

Acknowledging Realities, Not Excuses

It is important to acknowledge the genuine constraints. Many African tech companies, particularly startups, are indeed navigating challenging economic landscapes where every shilling counts. Implementing comprehensive mental health programs can seem like a daunting expense when survival itself is a daily concern. The argument of "resource constraints," while valid, risks becoming an excuse if mental health is not viewed as a critical investment in human capital—an investment that yields tangible returns in productivity, innovation, and retention.

Furthermore, mental well-being is a multifaceted societal challenge, influenced by factors far beyond the workplace, including the overall strength of national healthcare systems, which in many African countries are underfunded and under-resourced. The workplace cannot solve everything, but it has a profound impact and a significant role to play.

Encouragingly, awareness is growing. The rise of African mental health tech startups is a testament to this, even as it underscores a service gap that traditional employers, including tech companies themselves, have yet to fully address. These platforms offer innovative ways to bypass stigma and access barriers, yet their reach can be limited by the same infrastructural challenges—like internet accessibility or data costs—that contribute to workplace stress.

Debugging Our Approach: A Call for Action in African Tech

As we embrace the 2025 Mental Health Awareness Month theme, the well-being of Africa's tech talent must become a strategic priority. This is not merely an HR concern; it is fundamental to fostering sustainable innovation, driving economic growth, and ensuring our continent's technological advancement is both equitable and human-centered.

For tech leaders and companies in Uganda and across Africa, the first step is to cultivate cultures of openness and psychological safety. This means leadership openly discussing mental health, destigmatizing help-seeking, and establishing safe, confidential channels for employees to voice concerns without fear of reprisal. For those of us in AI/ML, this safety must extend to discussing the ethical quandaries and potential societal impacts of our work without fear of being penalized for raising difficult questions.

Secondly, invest in culturally-attuned, accessible mental health resources. This means moving beyond generic offerings to partner with local mental health professionals and organizations that understand the specific cultural contexts. Companies like Mindverse Uganda already offer workplace workshops. Services leveraging principles from successful local models like Uganda's TSP, which emphasizes cultural sensitivity, should be explored. Therapy benefits, flexible mental health leave policies, and confidential counselling should be the norm, not the exception.

Thirdly, design for well-being within work environments and project workflows. This includes promoting manageable workloads, offering flexible work arrangements where practical, ensuring clear role expectations, and embedding ethical AI development frameworks that explicitly consider developer well-being, especially when projects involve sensitive data or carry high societal impact.

Finally, empower managers and team leads with mental health literacy. Training them to recognize signs of distress, engage in supportive conversations, and guide team members to appropriate resources is crucial for early intervention.

As tech professionals, we also have a role. We must foster peer support networks and advocate responsibly for better conditions within our workplaces. Sharing our experiences, even anonymously, can build collective awareness and drive change. Crucially, we must prioritize practicing self-care and utilizing available resources, however limited they may seem. As ML engineers, we are uniquely positioned to not only highlight these issues but also to contribute to designing tech-enabled, data-informed mental health solutions for our workplaces and the wider community.

For policymakers and ecosystem enablers—governments, investors, and tech hubs—the call is to integrate mental health into tech development agendas. This means recognizing that a healthy workforce is a productive and innovative workforce. Investors can play a powerful role by incentivizing and supporting best practices, perhaps by including employee well-being metrics in their investment criteria or incubator programs.

Who cares about the mental health of Africa's tech talent? The answer must be a resounding "We do." The vision is an African tech ecosystem that is not only a global hub of innovation but also a beacon of human-centered progress. An ecosystem where the brilliant minds building the future, like Aisha, feel supported, valued, and mentally sound. Achieving this requires a collective commitment—from the boardroom to the individual developer, from policymakers to investors. Only then can we ensure that Africa's tech revolution truly benefits all.

Google has unveiled Gemma 3, the latest iteration of its open AI models, featuring four versions: gemma-3-1b-it, gemma-3-4b-it, gemma-3-12b-it, and gemma-3-27b-it.

The gemma-3-1b-it model is limited to text-only input, supports English exclusively, and comes with a 32k context length. Due to its lack of multilingual capabilities, it is unsuitable for a Luganda inference.

In contrast, the gemma-3-4b-it, gemma-3-12b-it, and gemma-3-27b-it models support both text and image input, recognize 140+ languages, and offer an extended 128k context length, making them far better suited for multilingual tasks.

For this specific task, we are using gemma-3-4b-it due to its balance of performance and efficiency.

Accessing Gemma 3 models

Before using Gemma 3 for the first time, you must request access to the model through Hugging Face by completing the following steps:

1. Log in to Hugging Face, or create a new Hugging Face account if you don't already have one.

2. Go to the Gemma 3 model card to get access to the model.

3. Complete the consent form and accept the terms and conditions.

To generate a Hugging Face token, open your Settings page in Hugging Face, choose Access Tokens option in the left pane and click New token. In the next window that appears, give a name to your token and choose the type as Write to get the write access.

Then, in Colab, select Secrets (🔑) in the left pane and add your Hugging Face token. Store your Hugging Face token under the name HF_TOKEN.

Select the runtime

To complete this tutorial, you'll need to have a Colab runtime with sufficient resources to load the Gemma 3 model. In this case, a T4/L4 GPU would be needed to load the model weights.

1. In the upper-right of the Colab window, click the dropdown menu.

2. Select Change runtime type.

3. Under Hardware accelerator, select T4 or L4.

Install Transformers

!pip install git+https://github.com/huggingface/transformers@v4.49.0-Gemma-3

Import libraries and dependencies

from transformers import AutoProcessor, Gemma3ForConditionalGeneration
from PIL import Image
import cv2
from IPython.display import Markdown, HTML
from base64 import b64encode
import requests
import torch

Choose the Gemma 3 model variant to use

Gemma 3 is available in four sizes, each supporting different features:

• gemma-3-1b-it

  • Supports only text input and English language

  • 32k context length

• gemma-3-4b-it, gemma-3-12b-it, gemma-3-27b-it

  • Supports both text and image input

  • Supports 140+ languages

  • 128k context length

model_name = 'gemma-3-4b-it' #We are using 4b
model_id = f"google/{model_name}"

model = Gemma3ForConditionalGeneration.from_pretrained(
    model_id, device_map="auto", torch_dtype=torch.bfloat16,
).eval()

processor = AutoProcessor.from_pretrained(model_id)

Define helper functions

• resize_image: Resizes the input images to n x n pixels, ensuring the aspect ratio is preserved.

• get_model_response: Send a text prompt and an image to the model, and retrieve the model's response.

• extract_frames: Extracts a specified number of evenly spaced frames from a video file along with their timestamps.

• show_video: Embeds and displays a video in an HTML5 player.

def resize_image(image_path):
    img = Image.open(image_path)

    target_width, target_height = 640, 640
    # Calculate the target size (maximum width and height).
    if target_width and target_height:
        max_size = (target_width, target_height)
    elif target_width:
        max_size = (target_width, img.height)
    elif target_height:
        max_size = (img.width, target_height)

    img.thumbnail(max_size)

    return img


def get_model_response(img: Image, prompt: str, model, processor):
    # Prepare the messages for the model.
    messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": "You are a helpful assistant. Reply only with the answer to the question asked in Luganda language only, and avoid using additional text in your response like 'here's the answer'."}]
        },
        {
            "role": "user",
            "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": prompt}
            ]
        }
    ]

    # Tokenize inputs and prepare for the model.
    inputs = processor.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=True,
        return_dict=True, return_tensors="pt"
    ).to(model.device, dtype=torch.bfloat16)

    input_len = inputs["input_ids"].shape[-1]

    # Generate response from the model.
    with torch.inference_mode():
        generation = model.generate(**inputs, max_new_tokens=100, do_sample=False)
        generation = generation[0][input_len:]

    # Decode the response.
    response = processor.decode(generation, skip_special_tokens=True)
    return response


def extract_frames(video_path, num_frames):
    """
    The function is adapted from:
    https://github.com/merveenoyan/smol-vision/blob/main/Gemma_3_for_Video_Understanding.ipynb
    """
    cap = cv2.VideoCapture(video_path)

    if not cap.isOpened():
        print("Error: Could not open video file.")
        return []

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    # Calculate the step size to evenly distribute frames across the video.
    step = total_frames // num_frames
    frames = []

    for i in range(num_frames):
        frame_idx = i * step
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
        ret, frame = cap.read()
        if not ret:
            break
        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        timestamp = round(frame_idx / fps, 2)
        frames.append((img, timestamp))

    cap.release()
    return frames


def show_video(video_path, video_width = 600):
  video_file = open(video_path, "r+b").read()
  video_url = f"data:video/mp4;base64,{b64encode(video_file).decode()}"
  video_html = f"""<video width={video_width} controls><source src="{video_url}"></video>"""
  return HTML(video_html)

Run an inference on images

Fetch some sample images for inferencing.

!wget https://raw.githubusercontent.com/wkambale/Luganda-Inference-on-Gemma-3/main/assets/image_1.jpg -O /content/image_1.jpg
!wget https://raw.githubusercontent.com/wkambale/Luganda-Inference-on-Gemma-3/main/assets/image_2.jpg -O /content/image_2.jpg
!wget https://raw.githubusercontent.com/wkambale/Luganda-Inference-on-Gemma-3/main/assets/image_3.jpg -O /content/image_3.jpg
!wget https://raw.githubusercontent.com/wkambale/Luganda-Inference-on-Gemma-3/main/assets/image_4.png -O /content/image_4.png

• Image 1 Credit: The Pearl

Task 1: Describe an image

The prompt is in Luganda language which translates to: "Describe the image."

image_file = 'image_1.jpg'
prompt = "Nnyonnyola emmeere eri mu kifaananyi."


img = resize_image(image_file)
display(img)
response = get_model_response(img, prompt, model, processor)
display(Markdown(response))

Response:

Omukoyogo.

Example 2: Identify a landmark

The prompt is in Luganda language which translates to: "Identify the famous landmark and location"

• Image 2: Ali Zali

image_file = 'image_2.jpg'
prompt = "Londoola ekifo kino ekimanyiddwa ennyo nne w'ekisangibwa."

img = resize_image(image_file)
display(img)
response = get_model_response(img, prompt, model, processor)
display(Markdown(response))

Response:

Ebibuga.

• Image 3: The Tower Post

image_file = 'image_3.jpg'
prompt = "Londoola ekifo kino ekimanyiddwa ennyo nne w'ekisangibwa."

img = resize_image(image_file)
display(img)
response = get_model_response(img, prompt, model, processor)
display(Markdown(response))

Response:

Kampala Bbwalo.

Task 3: Mathematical Reasoning/Okulowooza mu Kubala

The prompt is in Luganda language which translates to: "What is the value of x?"

• Image: Nitin

from PIL import Image
from IPython.display import Markdown

image_file = 'image_4.png'
prompt = "Omuwendo gwa x gwe guliwa?"

img = resize_image(image_file)
display(img)
response = get_model_response(img, prompt, model, processor)
display(Markdown(response))

Response:

x = 3

Inference on videos

The video is a clip from "Why Uganda is the Pearl Of Africa!" shoot.

• Credits: Eunice Tess

• Source: YouTube

# Video file.
video_path = "video.mp4"

# No. of frames to be extracted from the video.
num_frames = 10

video_output = show_video(video_path, video_width=800)
display(video_output)

The prompt is in Luganda language which translates to: "Please summarize what is happening in this video"

video_frames = extract_frames(video_path, num_frames=num_frames)

messages = [
    {
        "role": "system",
        "content": [{"type": "text", "text": "You are a helpful assistant."}]
    },
    {
        "role": "user",
        "content": [{"type": "text", "text": "Nsaba mufunze ebigenda mu maaso mu katambi kano"}]
    }
]


# Add frames to the messages structure.
for frame_data in video_frames:
    img, timestamp = frame_data
    messages[1]["content"].append({"type": "text", "text": f"Frame at {timestamp} seconds:"})
    img.save(f"/content/frames/frame_{timestamp}.png")
    messages[1]["content"].append({"type": "image", "url": f"/content/frames/frame_{timestamp}.png"})


inputs = processor.apply_chat_template(
    messages, add_generation_prompt=True, tokenize=True,
    return_dict=True, return_tensors="pt"
).to(model.device)


input_length = inputs["input_ids"].shape[-1]

# Generate a response based on the inputs.
output = model.generate(**inputs, max_new_tokens=500, do_sample=False)
output = output[0][input_length:]
response = processor.decode(output, skip_special_tokens=True)

display(Markdown(response))

Response:

Okay, let's look at these images. Here's what I see in Luganda:

Frame 0.0: "Nyo mu maaso, ttiima nnyo. Twee nyo mu nsi y'omutwe, nti nyo mu maaso." (It's beautiful, very impressive. It's in a world of wonder, it's in the water.)
Frame 13.2: "Nyo mu maaso, ttiima nnyo. Twee nyo mu nsi y'omutwe, nti nyo mu maaso." (It's beautiful, very impressive. It's in a world of wonder, it's in the water.)
Frame 26.4: "Nyo mu maaso, ttiima nnyo. Twee nyo mu nsi y'omutwe, nti nyo mu maaso." (It's beautiful, very impressive. It's in a world of wonder, it's in the water.)
Frame 39.6: "Nyo mu maaso, ttiima nnyo. Twee nyo mu nsi y'omutwe, nti nyo mu maaso." (It's beautiful, very impressive. It's in a world of wonder, it's in the water.)
Frame 52.8: "Nyo mu maaso, ttiima nnyo. Twee nyo mu nsi y'omutwe, nti nyo mu maaso." (It's beautiful, very impressive. It's in a world of wonder, it's in the water.)
Frame 66.0: "Nyo mu maaso, ttiima nnyo. Twee nyo mu nsi y'omutwe, nti nyo mu maaso." (It's beautiful, very impressive. It's in a world of wonder, it's in the water.)
Frame 79.2: "Nyo mu maaso, ttiima nnyo. Twee nyo mu nsi y'omutwe, nti nyo mu maaso."

Deductions

The outputs above reveal a significant limitation: despite Gemma 3 models boasting multilingual capabilities across 140+ languages, they still struggle to handle vision tasks (images and videos) in Luganda effectively.

This demonstration underscores the urgent need for:

1. More research into optimizing AI models for low-resource languages like Luganda.

2. Expanding datasets with high-quality, Luganda-specific image and video annotations.

3. Training foundational models that natively understand Luganda in multimodal contexts.

Without these critical steps, AI-powered vision systems will continue to exclude Luganda and other underrepresented languages from advancements in multimodal AI.

Resources

Here is the notebook.

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:

1. Automatic Vectorization – With functions like vmap, JAX makes it easy to apply operations over large batches of data without writing explicit loops.

2. Efficient Autograd – JAX provides automatic differentiation using grad, which simplifies training deep learning models.

3. XLA Compilation – Just-In-Time (JIT) compilation speeds up execution by compiling computation graphs for efficient hardware utilization.

4. 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  # Image processing in JAX

Verify GPU Access

JAX runs computations on CPUs, GPUs, and TPUs seamlessly. To check if your machine has a GPU or TPU available:

# Get available devices
print("Available Devices:", jax.devices())

# Check if GPU is available
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:

• A training dataset (train_data)

• A test dataset (test_data)

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) # k = number of times the rotation is applied

    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):
    # Rescale and resize all datasets.
    ds = ds.map(lambda x, y: (resize_and_rescale(x), y), num_parallel_calls=AUTOTUNE)

    if shuffle:
        ds = ds.shuffle(1000)

    # Use buffered prefetching on all datasets.
    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  # Number of output classes

    @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))  # Flatten feature maps
        x = nn.Dense(features=128)(x)
        x = nn.relu(x)

        x = nn.Dense(features=self.num_classes)(x)  # Output layer
        return x

1. 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.

2. 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.

3. Third Convolutional Layer

   • Increases the number of channels to 128 while keeping the 3×3 kernel size.

   • Applies ReLU activation and another max pooling step.

4. 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

# Set up PRNG key
key = jax.random.PRNGKey(0)

# Define input shape (batch_size, height, width, channels)
input_shape = (1, 32, 32, 3)  # Example for a 32x32 RGB image

# Initialize model
model = CNN(num_classes=10)  # Assuming 10 output classes
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  # No additional attributes needed for now

# Define the optimizer
learning_rate = 0.001
optimizer = optax.adam(learning_rate)

# Initialize the training state
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

    # Compute gradients
    grads = jax.grad(loss_fn)(state.params)

    # Update model state
    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.

A key mandate of these communities is to organize regular meetups, workshops, summits, and conferences. These gatherings are typically hands-on, bringing together expert developers to share skills, knowledge, and technological advancements with attendees through interactive workshops, code labs, and hackathons. In Uganda, groups such as GDG Cloud Kampala and GDG Cloud Mbarara have consistently held technical events like DevFests, while the Python community has successfully organized PyCon Uganda, now heading into its third edition.

However, alongside the rise of these developer-led events, corporate entities have increasingly taken an interest in organizing technology-related summits and conferences. This trend, while seemingly positive, presents a growing concern: many of these corporate-backed events are significantly lacking in technical depth, especially in the crucial domain of artificial intelligence (AI).

The Corporate Takeover and the Technical Deficit

These corporate-led summits and conferences often secure substantial funding, sometimes from government entities and corporate sponsors eager to see their logos on billboards and banners. This kind of financial backing is something developer-led community events can only dream of. Having organized conferences like Google I/O Extended and DevFest Mbarara for three years—each attracting close to 200 students and young professionals—I have personally experienced the uphill battle of securing sponsorship. The struggle to find support is a familiar nightmare for many tech community organizers, who must often rely on volunteers and minimal budgets to pull off impactful events.

Yet, despite their financial strength, corporate-led events frequently fall short of delivering meaningful technical content. A crucial weakness is the lack of hands-on sessions where attendees can actively engage in building, deploying, and experimenting with AI models and other emerging technologies. In many cases, not a single AI model is built or deployed throughout these high-profile summits. The reason? The invited speakers are usually corporate executives, heads of IT departments, and business strategists who, while experienced in management, often lack deep familiarity with emerging AI technologies.

The AI Knowledge Gap at Corporate Summits

To be clear, this is not to undermine the experience and expertise of IT heads in various companies and entities. Many of them have led digital transformation projects, implemented enterprise systems, and managed IT infrastructure at scale. However, AI is a different ballgame. The ability to discuss generative AI tools like OpenAI’s ChatGPT, Google’s Gemini, xAI’s Grok, or DeepSeek’s R1 is one thing; understanding how to fine-tune or integrate large language models (LLMs) such as Gemma 2, LLaMA 3, or OpenAI’s GPT series to solve local problems another.

Few, if any, of these corporate speakers can claim hands-on experience in fine-tuning models, optimizing neural networks, or deploying AI systems into production. This expertise gap means that rather than meaningful, technical AI discussions, these summits often feature surface-level talks filled with buzzwords and high-level strategies that lack practical application. Meanwhile, developers seeking actual skills in AI model development and deployment leave these events empty-handed.

Contrast this with many corporate-led AI summits that focus solely on panel discussions and keynote addresses. While such formats are useful for high-level industry insights, they do little to equip attendees with the hands-on skills necessary for AI development. Without technical immersion, these summits risk becoming echo chambers where the same concepts are discussed repeatedly without any tangible output.

The Danger of Mimicking Bureaucratic Inefficiency

We frequently criticize governments for their endless boardroom meetings and benchmarking trips that yield little practical implementation. It would be a grave mistake to let this inefficiency seep into the developer community. If AI is the future, then our engagement with it must be hands-on. We cannot afford to merely discuss AI trends at conferences while leaving the actual coding and model development to others. We must actively build and deploy AI models, integrate them into real-world applications, and refine them through continuous experimentation.

Shifting the Focus: How AI Events Can Be More Impactful

To ensure that our AI summits and conferences are truly impactful, organizers should:

• Prioritize Hands-On Workshops – Every AI conference should include code labs where participants can build and deploy AI models in real time. For example, workshops on fine-tuning LLMs, training custom vision models, or implementing AI in cloud environments should be staple sessions.

• Feature Technical Speakers – Rather than filling panels with corporate executives, AI summits should prioritize experts with hands-on experience in AI research and development. This means bringing in data scientists, machine learning engineers, and AI practitioners who can demonstrate real-world applications.

• Incorporate Hackathons and AI Challenges – Practical AI challenges, such as hackathons, Kaggle-style competitions, or live coding challenges, should be an integral part of any AI summit.

• Encourage Open-Source Contributions – Conferences should foster contributions to AI open-source projects, enabling attendees to leave with more than just theoretical knowledge.

• Partner with Developer Communities – Corporate entities should work with developer communities to deliver the hands-on approach to their summits. These communities have the target audience and the experience in organizing these events.

• Develop Local AI Talent – AI conferences should facilitate mentorship programs that connect experienced AI practitioners with budding developers to ensure continuous learning beyond the event itself.

The Call to Action: Time to Get Technical

If we are serious about competing in the AI race, we must radically rethink the structure of our summits and conferences. The AI field evolves at a blistering pace, and merely talking about it will leave us perpetually playing catch-up. To secure our place in the global AI ecosystem, we must do the actual work: writing the code, training the models, deploying real-world AI solutions, and iterating on them.

Uganda’s developer community has shown immense potential, but this potential must be nurtured with the right kind of engagement. Technical summits, practical workshops, and real-world AI applications should be our focus. Anything less than this risks turning our AI discourse into an empty parade of slogans and missed opportunities.

We must choose: do we want to be passive spectators in the AI revolution, or do we want to be active builders shaping the future? The answer lies in how we structure our conferences today.

…

Originally published in The Independent - Uganda here.

In this article, we explore how to deploy machine learning models in a scalable and efficient manner using TensorFlow Serving. TensorFlow Serving is a flexible, high-performance serving system designed for production environments, enabling you to serve your machine learning models to a large number of clients efficiently. We will cover the basics of TensorFlow Serving, how to set it up, how to serve models, and best practices for scaling your deployment.

What is TensorFlow Serving?

TensorFlow Serving is an open-source serving system specifically designed for deploying machine learning models in production environments. It allows you to serve multiple models or multiple versions of the same model simultaneously, and it can be easily integrated with TensorFlow models.

Why Use TensorFlow Serving?

Scalability: TensorFlow Serving is designed to handle high-throughput predictions, making it suitable for large-scale deployments.

Flexibility: It supports multiple models and versions, allowing for easy model management.

Efficiency: TensorFlow Serving is optimized for performance, with low latency and high throughput.

Setting Up TensorFlow Serving

Installation

TensorFlow Serving can be installed on various platforms, including Linux, macOS, and Windows. However, the most common way to get TensorFlow Serving up and running is through Docker, which simplifies the process and ensures a consistent environment.

Installing TensorFlow Serving on Linux

If you prefer to install TensorFlow Serving directly on your system, you can follow these steps:

echo "deb [arch=amd64] http://storage.googleapis.com/tensorflow-serving-apt stable tensorflow-model-server tensorflow-model-server-universal" | sudo tee /etc/apt/sources.list.d/tensorflow-serving.list > /dev/null
curl -fsSL https://storage.googleapis.com/tensorflow-serving-apt/tensorflow-serving.release.pub.gpg | sudo apt-key add -
sudo apt-get update && sudo apt-get install tensorflow-model-server

Verify Installation

tensorflow_model_server --version

Docker Setup

Using Docker is the recommended way to set up TensorFlow Serving, as it provides a consistent environment across different platforms.

1. Pull the TensorFlow Serving Docker Image

    docker pull tensorflow/serving
   

2. Run the Docker Container

    docker run -p 8501:8501 --name=tf_serving \
    --mount type=bind,source=$(pwd)/model,target=/models/model_name \
    -e MODEL_NAME=model_name -t tensorflow/serving
   

This command starts a TensorFlow Serving container, serving the model located in the model directory.

Serving a TensorFlow Model

Exporting a TensorFlow Model

Before serving a model, you need to export it in a format that TensorFlow Serving can understand. Typically, TensorFlow models are saved in the SavedModel format.

import tensorflow as tf

# Define a simple model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(32,)),
    tf.keras.layers.Dense(1)
])

# Save the model
model.save('/path/to/exported_model')

Loading the Model into TensorFlow Serving

With the model saved in the correct format, you can load it into TensorFlow Serving by pointing the server to the directory containing the exported model.

docker run -p 8501:8501 --name=tf_serving \
--mount type=bind,source=/path/to/exported_model,target=/models/my_model \
-e MODEL_NAME=my_model -t tensorflow/serving

Making Predictions via REST API

TensorFlow Serving provides a REST API to interact with your models. You can make predictions by sending HTTP POST requests.

curl -d '{"instances": [[1.0, 2.0, 5.0, 1.0, 2.0, 5.0]]}' \
  -X POST http://localhost:8501/v1/models/my_model:predict

This request sends a JSON payload containing the input data and receives the model's predictions as a response.

Scaling TensorFlow Serving

Horizontal and Vertical Scaling

• Horizontal Scaling: Involves adding more instances of TensorFlow Serving, distributing the load across multiple servers. This can be achieved using container orchestration platforms like Kubernetes.

• Vertical Scaling: Involves increasing the resources (CPU, memory) of a single TensorFlow Serving instance. This can be done by allocating more resources to the Docker container.

Load Balancing

Load balancing is crucial for handling large volumes of requests efficiently. You can use a load balancer to distribute incoming requests across multiple TensorFlow Serving instances.

Monitoring and Logging

Monitoring and logging are essential for understanding the performance of your TensorFlow Serving deployment. TensorFlow Serving integrates well with monitoring tools like Prometheus and Grafana.

Example Prometheus configuration:

global:
  scrape_interval: 15s

scrape_configs:
  - job_name: 'tensorflow_serving'
    static_configs:
      - targets: ['localhost:8501']

This configuration will scrape metrics from your TensorFlow Serving instance every 15 seconds.

Advanced Features of TensorFlow Serving

Model Versioning

TensorFlow Serving supports serving multiple versions of the same model. You can specify which version to serve using the --model_version_policy flag.

docker run -p 8501:8501 --name=tf_serving \
--mount type=bind,source=/path/to/exported_model,target=/models/my_model \
-e MODEL_NAME=my_model -e MODEL_VERSION_POLICY="latest" -t tensorflow/serving

Batch Prediction

TensorFlow Serving can be configured to perform batch predictions, which can significantly improve performance for high-throughput scenarios.

docker run -p 8501:8501 --name=tf_serving \
--mount type=bind,source=/path/to/exported_model,target=/models/my_model \
-e MODEL_NAME=my_model -e TF_SERVING_BATCHING_PARAMETERS_FILE="/path/to/batching_parameters" \
-t tensorflow/serving

Customizing TensorFlow Serving

You can customize TensorFlow Serving by adding custom code for preprocessing, postprocessing, or integrating with other systems.

Example of a custom model handler:

import tensorflow as tf
from tensorflow_serving.apis import predict_pb2, prediction_service_pb2_grpc

class CustomModelHandler:
    def __init__(self, model_path):
        self.model = tf.keras.models.load_model(model_path)

    def predict(self, request: predict_pb2.PredictRequest) -> predict_pb2.PredictResponse:
        # Custom preprocessing
        inputs = request.inputs['input_tensor'].numpy()

        # Model prediction
        predictions = self.model.predict(inputs)

        # Custom postprocessing
        response = predict_pb2.PredictResponse()
        response.outputs['output_tensor'].CopyFrom(tf.make_tensor_proto(predictions))
        return response

Best Practices

Security Considerations

• Authentication: Implement authentication mechanisms to ensure that only authorized clients can access your models.

• Encryption: Use TLS to encrypt data in transit between clients and TensorFlow Serving.

Resource Management

• CPU and Memory Limits: Set appropriate limits on CPU and memory usage to prevent resource exhaustion.

• Autoscaling: Use autoscaling to dynamically adjust the number of TensorFlow Serving instances based on demand.

Optimizing Performance

• Model Optimization: Optimize your model using techniques like quantization to reduce latency.

• Caching: Implement caching mechanisms to store frequently requested predictions, reducing the load on your TensorFlow Serving instance.

Conclusion

TensorFlow Serving provides a powerful and flexible solution for serving machine learning models in production environments. By following the steps outlined in this tutorial, you can set up a scalable TensorFlow Serving deployment that is capable of handling large volumes of requests efficiently. With advanced features like model versioning, batch prediction, and custom handlers, TensorFlow Serving can be tailored to meet the specific needs of your application.

What is TFX?

TFX is a production-ready ML platform designed to help you build, deploy, and manage ML models. It consists of a set of libraries and tools that help automate and manage the ML lifecycle. TFX pipelines are portable and can run on various platforms, including Apache Beam, Apache Airflow, and Kubeflow.

Key Features of TFX

• Scalability: TFX can handle large-scale data processing and training.

• Portability: Pipelines can run on different platforms and environments.

• Modularity: TFX components are designed to be modular, allowing you to customize and extend them as needed.

• Production-Ready: TFX is built with production deployment in mind, ensuring reliability and robustness.

TFX Components

TFX pipelines are composed of several components, each responsible for a specific part of the ML lifecycle. Here are the core components:

ExampleGen

ExampleGen is the first component in a TFX pipeline. It ingests and splits data into training and evaluation datasets. ExampleGen supports various data sources, such as CSV files, TFRecord files, and BigQuery.

from tfx.components import CsvExampleGen

example_gen = CsvExampleGen(input_base='/content/iris.csv')

StatisticsGen

StatisticsGen computes statistics over the data for data visualization and validation. It generates statistics using TensorFlow Data Validation (TFDV).

from tfx.components import StatisticsGen

statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])

SchemaGen

SchemaGen generates a schema for the data based on the statistics computed by StatisticsGen. The schema includes information about the data types, domains, and constraints.

from tfx.components import SchemaGen

schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'])

ExampleValidator

ExampleValidator detects anomalies in the data by comparing the data against the schema generated by SchemaGen. It helps ensure data quality.

from tfx.components import ExampleValidator

example_validator = ExampleValidator(statistics=statistics_gen.outputs['statistics'], schema=schema_gen.outputs['schema'])

Transform

Transform performs feature engineering and data transformation using TensorFlow Transform (TFT). It preprocesses the data for model training and serving.

from tfx.components import Transform

def preprocessing_fn(inputs):
    outputs = {
        'sepal_length': Transform.scale_to_z_score(inputs['sepal.length']),
        'sepal_width': Transform.scale_to_z_score(inputs['sepal.width']),
        'petal_length': Transform.scale_to_z_score(inputs['petal.ength']),
        'petal_width': Transform.scale_to_z_score(inputs['petal.width']),
        'species': inputs['variety']
    }
    return outputs

transform = Transform(
    examples=example_gen.outputs['examples'],
    schema=schema_gen.outputs['schema'],
    module_file='/content/preprocessing.py'
)

Trainer

Trainer trains an ML model using the preprocessed data. It supports various frameworks, including TensorFlow, Keras, and Estimator.

from tfx.components import Trainer
from tfx.proto import trainer_pb2

trainer = Trainer(
    module_file='/content/trainer_module.py',
    examples=transform.outputs['transformed_examples'],
    schema=schema_gen.outputs['schema'],
    transform_graph=transform.outputs['transform_graph'],
    train_args=trainer_pb2.TrainArgs(num_steps=1000),
    eval_args=trainer_pb2.EvalArgs(num_steps=100)
)

Evaluator

Evaluator evaluates the trained model using TensorFlow Model Analysis (TFMA). It helps in validating and comparing different models.

import tensorflow_model_analysis as tfma

eval_config = tfma.EvalConfig(
    slicing_specs=[tfma.SlicingSpec()],
    metrics_specs=[
        tfma.MetricsSpec(
            metrics=[
                tfma.MetricConfig(class_name='SparseCategoricalAccuracy')
            ]
        )
    ]
)

evaluator = Evaluator(
    examples=example_gen.outputs['examples'],
    model=trainer.outputs['model'],
    eval_config=eval_config
)

Pusher

Pusher deploys the trained model to a serving infrastructure. It ensures that the model meets certain criteria before pushing it to production.

from tfx.components import Pusher
from tfx.proto import pusher_pb2

pusher = Pusher(
    model=trainer.outputs['model'],
    model_blessing=evaluator.outputs['blessing'],
    push_destination=pusher_pb2.PushDestination(
        filesystem=pusher_pb2.PushDestination.Filesystem(
            base_directory='/content/model'
        )
    )
)

Setting Up TFX

Before building a TFX pipeline, it's essential to set up the environment. This involves installing the necessary packages and configuring the runtime environment.

Installing TFX: TFX can be installed via pip. The installation includes all the required libraries and dependencies for running TFX components.

pip install tfx

Configuring the Environment: Setting up the environment involves configuring paths for data, pipelines, and model artifacts. This ensures that all components can access the necessary resources and save outputs in the correct locations.

Building a Simple TFX Pipeline

To illustrate the capabilities of TFX, we will build a pipeline using the Iris dataset, a well-known dataset for classification tasks. The Iris dataset contains 150 samples of iris flowers, each with four features (sepal length, sepal width, petal length, petal width) and a class label (species).

Data Ingestion

The first step in the TFX pipeline is to ingest the Iris dataset using ExampleGen. This component reads the dataset, splits it into training and evaluation sets, and converts it into the TFX internal format.

CSV ExampleGen: For the Iris dataset, we use the CsvExampleGen component, which ingests data from CSV files. It automatically splits the data into training and evaluation sets based on a specified ratio.

from tfx.components import CsvExampleGen

example_gen = CsvExampleGen(input_base='/content/iris.csv')

Data Statistics

StatisticsGen computes descriptive statistics for the dataset, providing insights into data distributions and detecting anomalies. It uses TensorFlow Data Validation (TFDV) to generate statistics such as mean, median, and standard deviation for each feature.

from tfx.components import StatisticsGen

statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])

Importance: Understanding the distribution of data is crucial for identifying potential issues and ensuring that the data is suitable for training. StatisticsGen helps detect anomalies such as outliers and missing values, which can affect model performance.

Schema Generation

Based on the statistics computed by StatisticsGen, SchemaGen generates a schema for the dataset. The schema includes information about feature types, value ranges, and presence constraints, serving as a blueprint for data validation and transformation.

from tfx.components import SchemaGen

schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'])

Definition: The schema defines the expected structure of the data, including feature types (numeric, categorical, etc.), value ranges, and constraints (e.g., required features). This information is critical for ensuring data consistency and preparing it for model training.

Data Validation

ExampleValidator validates the dataset against the schema, identifying anomalies and missing values. It ensures that the data adheres to the expected format, which is essential for training reliable models.

from tfx.components import ExampleValidator

example_validator = ExampleValidator(statistics=statistics_gen.outputs['statistics'], schema=schema_gen.outputs['schema'])

Anomaly Detection: ExampleValidator detects anomalies such as outliers, missing values, and unexpected feature values. These issues can affect model performance and lead to unreliable predictions, making data validation a crucial step in the pipeline.

Data Transformation

Transform performs feature engineering and data preprocessing using TensorFlow Transform (TFT). It applies transformations such as scaling, normalization, and encoding, preparing the data for model training.

from tfx.components import Transform

def preprocessing_fn(inputs):
    # Normalize the numeric features
    outputs = {
        'sepal_length': Transform.scale_to_z_score(inputs['sepal.length']),
        'sepal_width': Transform.scale_to_z_score(inputs['sepal.width']),
        'petal_length': Transform.scale_to_z_score(inputs['petal.ength']),
        'petal_width': Transform.scale_to_z_score(inputs['petal.width']),
        'species': inputs['variety']
    }
    return outputs

transform = Transform(
    examples=example_gen.outputs['examples'],
    schema=schema_gen.outputs['schema'],
    module_file='/content/preprocessing.py'
)

Preprocessing Functions: Transform defines preprocessing functions that apply transformations to the raw data. These functions can include operations such as scaling numerical features, encoding categorical features, and generating new features based on existing ones.

Model Training

The Trainer component trains the ML model using the transformed data. It leverages TensorFlow's capabilities to define, train, and evaluate models.

import tensorflow as tf
from tfx import v1 as tfxio

def _input_fn(file_pattern, data_accessor, schema, batch_size=200):
    raw_data = data_accessor.tf_dataset_factory(
        file_pattern, tfxio.TensorFlowDatasetOptions(batch_size=batch_size), schema)
    transformed_data = raw_data.map(_parse_fn)
    return transformed_data

def _build_keras_model():
    model = tf.keras.Sequential([
        tf.keras.layers.InputLayer(input_shape=(4,)),
        tf.keras.layers.Dense(10, activation='relu'),
        tf.keras.layers.Dense(10, activation='relu'),
        tf.keras.layers.Dense(3, activation='softmax')
    ])
    model.compile(optimizer=tf.keras.optimizers.Adam(lr=0.001),
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    return model

def run_fn(fn_args):
    schema = tfx_bsl.public.tfxio_utils.get_tfx_schema_from_tensorflow_metadata_schema(
        fn_args.schema)
    train_dataset = _input_fn(fn_args.train_files, fn_args.data_accessor, schema, batch_size=200)
    eval_dataset = _input_fn(fn_args.eval_files, fn_args.data_accessor, schema, batch_size=200)

    model = _build_keras_model()
    model.fit(train_dataset, steps_per_epoch=fn_args.train_steps,
              validation_data=eval_dataset, validation_steps=fn_args.eval_steps)
    model.save(fn_args.serving_model_dir, save_format='tf')

Model Definition: Trainer uses a module file containing the model definition and training logic. This file defines the architecture of the model, the loss function, and the optimization algorithm. It also includes the training and evaluation steps, specifying the number of epochs, batch size, and evaluation metrics.

Trainer

Trainer trains an ML model using the preprocessed data. It supports various frameworks, including TensorFlow, Keras, and Estimator

from tfx.components import Trainer
from tfx.proto import trainer_pb2

trainer = Trainer(
    module_file='/content/trainer_module.py',
    examples=transform.outputs['transformed_examples'],
    schema=schema_gen.outputs['schema'],
    transform_graph=transform.outputs['transform_graph'],
    train_args=trainer_pb2.TrainArgs(num_steps=1000),
    eval_args=trainer_pb2.EvalArgs(num_steps=100)
)

Model Evaluation

Evaluator evaluates the trained model using TensorFlow Model Analysis (TFMA). It performs a detailed analysis of model performance, identifying potential issues and ensuring that the model meets the desired criteria before deployment.

import tensorflow_model_analysis as tfma

eval_config = tfma.EvalConfig(
    slicing_specs=[tfma.SlicingSpec()],
    metrics_specs=[
        tfma.MetricsSpec(
            metrics=[
                tfma.MetricConfig(class_name='SparseCategoricalAccuracy')
            ]
        )
    ]
)

evaluator = Evaluator(
    examples=example_gen.outputs['examples'],
    model=trainer.outputs['model'],
    eval_config=eval_config
)

Evaluation Configuration: Evaluator uses an evaluation configuration to specify the metrics and slicing specifications for model evaluation. Metrics such as accuracy, precision, and recall are used to assess model performance, while slicing specifications allow for analyzing performance across different subsets of the data.

Model Deployment

The final step in the TFX pipeline is to deploy the validated model using Pusher. This component ensures that only the best models are deployed to the serving infrastructure.

from tfx.components import Pusher
from tfx.proto import pusher_pb2

pusher = Pusher(
    model=trainer.outputs['model'],
    model_blessing=evaluator.outputs['blessing'],
    push_destination=pusher_pb2.PushDestination(
        filesystem=pusher_pb2.PushDestination.Filesystem(
            base_directory='/content/model'
        )
    )
)

Model Deployment: Pusher deploys the model to a specified serving infrastructure, such as TensorFlow Serving. It ensures that the model is production-ready and meets the desired performance criteria, facilitating continuous model improvement.

Running the Pipeline

The pipeline is executed using an orchestrator like the Local DAG Runner, which processes the data, trains the model, evaluates its performance, and deploys it if it meets the required criteria.

from tfx.orchestration import pipeline
from tfx.orchestration.local.local_dag_runner import LocalDagRunner

# Define the pipeline
pipeline = pipeline.Pipeline(
    pipeline_name='iris_pipeline',
    pipeline_root='/content/iris_pipeline',
    components=[example_gen, statistics_gen, schema_gen, example_validator, transform, trainer, evaluator, pusher],
    enable_cache=True,
    metadata_connection_config=None
)

# Run the pipeline
LocalDagRunner().run(pipeline)

After execution, the pipeline produces outputs such as transformed data, trained models, evaluation results, and metadata, all stored in a specified directory for further analysis and deployment.

Monitoring involves checking logs and reviewing evaluation metrics, while debugging includes inspecting artifacts and re-executing specific components to resolve any issues. This process ensures the pipeline runs smoothly, producing reliable and scalable machine learning models ready for deployment.

Conclusion

In this article, we introduced TensorFlow Extended (TFX) and its core components. We demonstrated how to set up a TFX environment, build a simple TFX pipeline using the Iris dataset, and run it. TFX provides a powerful and flexible framework for managing the end-to-end ML lifecycle, from data ingestion to model deployment. You should now have a solid foundation for building and deploying your own TFX pipelines.

TFX's modularity and scalability make it suitable for a wide range of ML applications, ensuring that you can build robust and production-ready ML systems. Happy experimenting with TFX!

Resources

GitHub Gist to the notebook

What is Distributed Training

Distributed training allows you to leverage multiple GPUs, TPUs, or even multiple machines to accelerate the training process of your machine learning models. TensorFlow's distributed training capabilities are built around the concept of a "distribution strategy," which specifies how computation is distributed across devices.

Types of Distributed Strategies

TensorFlow provides several strategies for distributed training, each suited for different scenarios and hardware configurations. Let's get into each strategy, including their use cases and advantages to help you get started.

MirroredStrategy

tf.distribute.MirroredStrategy is designed for synchronous training on multiple GPUs on a single machine. It replicates all of the model variables across the GPUs and then performs a synchronous update to keep them in sync.

MultiWorkerMirroredStrategy

tf.distribute.MultiWorkerMirroredStrategy extends MirroredStrategy to multiple machines. Each worker (machine) runs a replica of the model and synchronizes updates across all workers.

TPUStrategy

tf.distribute.TPUStrategy is used to train models on Google's TPUs. It is optimized for high-performance training and requires minimal code changes from GPU training.

ParameterServerStrategy

tf.distribute.experimental.ParameterServerStrategy is an asynchronous training strategy where the computation is divided between parameter servers and workers. Parameter servers store model parameters, and workers perform the computations.

Preparing the Data

Data preparation is a critical step in any machine learning workflow. For distributed training, the way you prepare and feed data to your model can significantly impact the training efficiency and performance. TensorFlow's tf.data API is a powerful tool for building input pipelines that can be easily integrated with distributed training.

Loading and Preprocessing Data

We will use the MNIST dataset, a classic dataset of handwritten digits. The dataset is available directly through TensorFlow, which makes loading and preprocessing straightforward.

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

x_train = x_train / 255.0
x_test = x_test / 255.0

x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

Creating TensorFlow Datasets

TensorFlow Datasets (tf.data.Dataset) provides a high-level API for creating and manipulating data pipelines. Using this API, we can create efficient input pipelines that are capable of feeding data to the model in a scalable and efficient manner.

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))

BATCH_SIZE = 64
train_dataset = train_dataset.shuffle(buffer_size=10000).batch(BATCH_SIZE)
test_dataset = test_dataset.batch(BATCH_SIZE)

Optimizing Data Pipelines

For distributed training, it’s important to ensure that the data pipeline does not become a bottleneck. TensorFlow provides several techniques to optimize data pipelines:

• Prefetching: Overlap the preprocessing and model execution of data.

• Caching: Cache data in memory to avoid redundant computations.

• Parallel Interleave: Read data from multiple files in parallel.

AUTOTUNE = tf.data.experimental.AUTOTUNE

train_dataset = train_dataset.cache()
train_dataset = train_dataset.prefetch(buffer_size=AUTOTUNE)
test_dataset = test_dataset.prefetch(buffer_size=AUTOTUNE)

Defining the Model

Defining a model in TensorFlow is typically done using the Keras API, which provides a simple and flexible way to build neural networks. Let's define a convolutional neural network (CNN) for the MNIST dataset.

Creating the Model

A CNN is well-suited for image classification tasks. Here, we'll create a simple CNN with two convolutional layers followed by pooling layers, a flattening layer, and two dense layers.

def create_model():
    model = tf.keras.Sequential([
        tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    return model

Compiling the Model

After defining the model, the next step is to compile it. Compilation involves specifying the optimizer, loss function, and metrics that the model should use during training.

model = create_model()
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

Model Summary

It’s always a good practice to print the model summary to understand the architecture and ensure that the model is correctly defined.

model.summary()

Configuring the Distributed Strategy

TensorFlow's distribution strategies allow you to run your training on multiple GPUs, TPUs, or even across multiple machines. This section explains how to set up and configure different distributed strategies.

MirroredStrategy

tf.distribute.MirroredStrategy is designed for synchronous training on multiple GPUs on a single machine. It replicates all model variables across the GPUs and then performs a synchronous update to keep them in sync.

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

MultiWorkerMirroredStrategy

tf.distribute.MultiWorkerMirroredStrategy extends MirroredStrategy to multiple machines. You need to configure the cluster spec and set the environment variables appropriately.

Setting Up Cluster Spec

cluster_spec = {
    'worker': ['worker1.example.com:2222', 'worker2.example.com:2222']
}

os.environ['TF_CONFIG'] = json.dumps({
    'cluster': cluster_spec,
    'task': {'type': 'worker', 'index': 0}
})

strategy = tf.distribute.MultiWorkerMirroredStrategy()

Training with MultiWorkerMirroredStrategy

with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    model.fit(train_dataset, epochs=5)
    model.evaluate(test_dataset)

TPUStrategy

tf.distribute.TPUStrategy is used to train models on Google's TPUs. It is optimized for high-performance training and requires minimal code changes from GPU training.

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='your-tpu-address')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)

with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    model.fit(train_dataset, epochs=5)
    model.evaluate(test_dataset)

ParameterServerStrategy

tf.distribute.experimental.ParameterServerStrategy is an asynchronous training strategy where the computation is divided between parameter servers and workers. Parameter servers store model parameters, and workers perform the computations.

cluster_spec = {
    'worker': ['worker1.example.com:2222', 'worker2.example.com:2222'],
    'ps': ['ps0.example.com:2222']
}

os.environ['TF_CONFIG'] = json.dumps({
    'cluster': cluster_spec,
    'task': {'type': 'worker', 'index': 0}

strategy = tf.distribute.experimental.ParameterServerStrategy()

Training with ParameterServerStrategy

with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    model.fit(train_dataset, epochs=5)
    model.evaluate(test_dataset)

Monitoring and Debugging

Monitoring and debugging distributed training can be challenging due to the complexity and scale of operations. TensorFlow provides several tools to help with this process, including TensorBoard, logging, and callbacks.

Using TensorBoard

TensorBoard is a powerful visualization tool that allows you to track and visualize metrics such as loss and accuracy during training. It can also display graphs, histograms, and other metrics to help you understand your model's behavior.

To use TensorBoard, you need to set up a TensorBoard callback during model training. This callback will log the metrics to a specified directory.

log_dir = "logs/"
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)

with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

model.fit(train_dataset, epochs=5, callbacks=[tensorboard_callback])
model.evaluate(test_dataset)

Launching TensorBoard

To launch TensorBoard, run the following command in your terminal:

tensorboard --logdir=logs/

This will start a local server where you can visualize the training metrics. Open your browser and navigate to http://localhost:6006/ to view the TensorBoard dashboard.

Using Logging

Logging is another useful tool for monitoring and debugging your training process. You can use Python’s built-in logging module to log messages and metrics during training.

logging.basicConfig(level=logging.INFO)
logging.info("Starting model training...")

with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

model.fit(train_dataset, epochs=5, callbacks=[tensorboard_callback])
model.evaluate(test_dataset)

logging.info("Model training completed.")

Using Callbacks

Callbacks are powerful tools that allow you to perform actions at various stages of the training process. TensorFlow provides several built-in callbacks, and you can also create custom callbacks to suit your needs.

Built-In Callbacks

TensorFlow includes several built-in callbacks, such as EarlyStopping, ModelCheckpoint, and ReduceLROnPlateau.

early_stopping_callback = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=3)
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath='model.h5', save_best_only=True)
reduce_lr_callback = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=2)

with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    model.fit(train_dataset, epochs=5, validation_data=test_dataset,
              callbacks=[tensorboard_callback, early_stopping_callback, model_checkpoint_callback, reduce_lr_callback])

Custom Callbacks

You can also create custom callbacks by subclassing tf.keras.callbacks.Callback.

class CustomCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        logging.info(f"Epoch {epoch} ended with loss: {logs['loss']} and accuracy: {logs['accuracy']}")

with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    model.fit(train_dataset, epochs=5, validation_data=test_dataset,
              callbacks=[tensorboard_callback, CustomCallback()])

Debugging with tf.debugging

TensorFlow also provides debugging tools in

the tf.debugging module to catch and diagnose issues during training. For example, you can use tf.debugging.assert_equal to ensure that tensors have expected values.

a = tf.constant(1)
b = tf.constant(2)

tf.debugging.assert_equal(a, b, message="Tensors are not equal")

Conclusion

Distributed training with TensorFlow can significantly accelerate the training process of your models by leveraging multiple devices and machines. This article covered the basics of setting up and running distributed training using various distribution strategies provided by TensorFlow. By understanding and utilizing these strategies, you can scale your machine learning workflows to handle larger datasets and more complex models efficiently.

Here is a summary of what we covered:

1. Introduction to Distributed Training: Understanding the need and benefits of distributed training.

2. Types of Distributed Strategies: Exploring different strategies like MirroredStrategy, MultiWorkerMirroredStrategy, TPUStrategy, and ParameterServerStrategy.

3. Preparing the Data: Loading and preprocessing the dataset.

4. Defining the Model: Creating a simple CNN model using TensorFlow's Keras API.

5. Configuring the Distributed Strategy: Setting up the appropriate distribution strategy for your training.

6. Monitoring and Debugging: Using TensorBoard to monitor and debug the training process.

With this knowledge, you are now equipped to start leveraging the power of distributed training to build and train more efficient and scalable machine learning models. Happy coding!

Done? Okay. Let's go...

Implementing Attention Mechanisms

Attention mechanisms have revolutionized the field of deep learning, particularly in natural language processing (NLP) and computer vision. They allow models to focus on specific parts of the input sequence or data, effectively improving the model's ability to capture dependencies and relationships.

Understanding Attention

Attention mechanisms work by assigning different weights to different parts of the input, allowing the model to focus on the most relevant parts. This is particularly useful in sequence-to-sequence tasks such as machine translation, where certain words in the input sequence may be more important than others for generating the output sequence.

Types of Attention

1. Self-Attention: Computes attention weights within the same sequence, allowing each element to focus on other elements in the sequence.

2. Cross-Attention: Computes attention weights between two different sequences, such as in encoder-decoder models.

Scaled Dot-Product Attention

The scaled dot-product attention mechanism is a common type of attention used in many models, including the Transformer. It involves three main components: queries (Q), keys (K), and values (V).

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Multi-Head Attention

Multi-head attention extends the concept of single attention by applying multiple attention mechanisms in parallel, allowing the model to focus on different parts of the input sequence simultaneously.

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, ..., \text{head}_h)W^O$$

where each head is an independent attention mechanism.

Implementing Attention in TensorFlow

Here's an example of implementing scaled dot-product attention and multi-head attention in TensorFlow.

Scaled Dot-Product Attention

import tensorflow as tf

def scaled_dot_product_attention(q, k, v, mask=None):
    matmul_qk = tf.matmul(q, k, transpose_b=True)
    dk = tf.cast(tf.shape(k)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    if mask is not None:
        scaled_attention_logits += (mask * -1e9)

    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
    output = tf.matmul(attention_weights, v)

    return output, attention_weights

Multi-Head Attention

class MultiHeadAttention(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model

        assert d_model % self.num_heads == 0

        self.depth = d_model // self.num_heads

        self.wq = tf.keras.layers.Dense(d_model)
        self.wk = tf.keras.layers.Dense(d_model)
        self.wv = tf.keras.layers.Dense(d_model)

        self.dense = tf.keras.layers.Dense(d_model)

    def split_heads(self, x, batch_size):
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, v, k, q, mask):
        batch_size = tf.shape(q)[0]

        q = self.wq(q)
        k = self.wk(k)
        v = self.wv(v)

        q = self.split_heads(q, batch_size)
        k = self.split_heads(k, batch_size)
        v = self.split_heads(v, batch_size)

        scaled_attention, attention_weights = scaled_dot_product_attention(q, k, v, mask)

        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])
        concat_attention = tf.reshape(scaled_attention, (batch_size, -1, self.d_model))

        output = self.dense(concat_attention)

        return output, attention_weights

Using Attention in a Transformer Model

The Transformer model relies heavily on attention mechanisms. Here's a brief overview of how attention is used in the Transformer architecture.

Transformer Encoder

The encoder consists of multiple layers, each containing a multi-head self-attention mechanism and a feed-forward neural network.

class EncoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super(EncoderLayer, self).__init__()
        self.mha = MultiHeadAttention(d_model, num_heads)
        self.ffn = tf.keras.Sequential([
            tf.keras.layers.Dense(dff, activation='relu'),
            tf.keras.layers.Dense(d_model)
        ])

        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)

    def call(self, x, training, mask):
        attn_output, _ = self.mha(x, x, x, mask)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(x + attn_output)

        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        out2 = self.layernorm2(out1 + ffn_output)

        return out2

Building Generative Models

Generative models are a class of machine learning models that learn to generate new data samples that resemble the training data. Two popular types of generative models are Variational Autoencoders (VAEs) and Generative Adversarial Networks (GANs).

Variational Autoencoders (VAEs)

Variational Autoencoders (VAEs) are probabilistic graphical models that aim to learn a latent representation of the data, which can then be used to generate new samples. VAEs consist of two main components: the encoder and the decoder.

Key Components of VAEs

1. Encoder: Maps the input data to a latent space, producing a mean and a variance for each dimension of the latent space.

2. Decoder: Maps the latent representation back to the data space, generating new samples that resemble the original data.

Implementing a VAE in TensorFlow

Here's an example of implementing a simple VAE for generating images.

import tensorflow as tf
from tensorflow.keras import layers, models

class Sampling(layers.Layer):
    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

latent_dim = 2
encoder_inputs = tf.keras.Input(shape=(28, 28, 1))
x = layers.Flatten()(encoder_inputs)
x = layers.Dense(512, activation='relu')(x)
z_mean = layers.Dense(latent_dim)(x)
z_log_var = layers.Dense(latent_dim)(x)
z = Sampling()([z_mean, z_log_var])
encoder = tf.keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")

decoder_inputs = tf.keras.Input(shape=(latent_dim,))
x = layers.Dense(512, activation='relu')(decoder_inputs)
x = layers.Dense(28 * 28 * 1, activation='sigmoid')(x)
decoder_outputs = layers.Reshape((28, 28, 1))(x)
decoder = tf.keras.Model(decoder_inputs, decoder_outputs, name="decoder")

class VAE(tf.keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def call(self, inputs):
        z_mean, z_log_var, z = self.encoder(inputs)
        reconstructed = self.decoder(z)
        kl_loss = -0.5 * tf.reduce_mean(z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1)
        self.add_loss(kl_loss)
        return reconstructed

vae = VAE(encoder, decoder)
vae.compile(optimizer='adam', loss='binary_crossentropy')

(x_train, _), (x_test, _) = tf.keras.datasets.mnist.load_data()
x_train = x_train.astype("float32") / 255.0
x_train = x_train.reshape(-1, 28, 28, 1)
x_test = x_test.astype("float32") / 255.0
x_test = x_test.reshape(-1, 28, 28, 1)

vae.fit(x_train, x_train, epochs=30, batch_size=128, validation_data=(x_test, x_test))

• Encoder: The encoder consists of a dense layer followed by two output layers: one for the mean and one for the log variance of the latent space.

• Decoder: The decoder maps the latent space back to the original data space.

• Sampling Layer: The sampling layer implements the reparameterization trick, which allows backpropagation through the stochastic latent space.

• VAE Model: The VAE model combines the encoder and decoder, adding the KL divergence loss to encourage the latent space to follow a standard normal distribution.

Generative Adversarial Networks (GANs)

Generative Adversarial Networks (GANs) consist of two neural networks: the generator and the discriminator. The generator learns to produce realistic data samples, while the discriminator learns to distinguish between real and generated samples. The two networks are trained in a competitive process.

Key Components of GANs

1. Generator: Takes random noise as input and generates data samples.

2. Discriminator: Takes data samples as input and classifies them as real or fake.

Implementing a GAN in TensorFlow

Here's an example of implementing a simple GAN for generating images.

import tensorflow as tf
from tensorflow.keras import layers, models

def build_generator():
    model = models.Sequential()
    model.add(layers.Dense(256, activation='relu', input_dim=100))
    model.add(layers.BatchNormalization())
    model.add(layers.Dense(512, activation='relu'))
    model.add(layers.BatchNormalization())
    model.add(layers.Dense(1024, activation='relu'))
    model.add(layers.BatchNormalization())
    model.add(layers.Dense(28 * 28 * 1, activation='tanh'))
    model.add(layers.Reshape((28, 28, 1)))
    return model

def build_discriminator():
    model = models.Sequential()
    model.add(layers.Flatten(input_shape=(28, 28, 1)))
    model.add(layers.Dense(512, activation='relu'))
    model.add(layers.Dense(256, activation='relu'))
    model.add(layers.Dense(1, activation='sigmoid'))
    return model

generator = build_generator()
discriminator = build_discriminator()
discriminator.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

gan_input = tf.keras.Input(shape=(100,))
generated_image = generator(gan_input)
discriminator.trainable = False
gan_output = discriminator(generated_image)
gan = tf.keras.Model(gan_input, gan_output)
gan.compile(optimizer='adam', loss='binary_crossentropy')

(x_train, _), (_, _) = tf.keras.datasets.mnist.load_data()
x_train = x_train.astype("float32") / 255.0
x_train = x_train.reshape(-1, 28, 28, 1)

import numpy as np

batch_size = 128
epochs = 10000
half_batch = batch_size // 2

for epoch in range(epochs):
    idx = np.random.randint(0, x_train.shape[0], half_batch)
    real_images = x_train[idx]
    noise = np.random.normal(0, 1, (half_batch, 100))
    fake_images = generator.predict(noise)

    d_loss_real = discriminator.train_on_batch(real_images, np.ones((half_batch, 1)))
    d_loss_fake = discriminator.train_on_batch(fake_images, np.zeros((half_batch, 1)))
    d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)

    noise = np.random.normal(0, 1, (batch_size, 100))
    valid_y = np.array([1] * batch_size)
    g_loss = gan.train_on_batch(noise, valid_y)

    if epoch % 100 == 0:
        print(f"{epoch} [D loss: {d_loss[0]} | D accuracy: {100*d_loss[1]}] [G loss: {g_loss}]")

• Generator: The generator network consists of dense layers followed by batch normalization and activation functions. It maps random noise to a data sample.

• Discriminator: The discriminator network consists of dense layers and activation functions. It classifies data samples as real or fake.

• Training Loop: The GAN is trained in a loop where the discriminator is trained on real and fake samples, followed by training the generator to produce samples that can fool the discriminator.

Hyperparameter Tuning and Model Evaluation

Hyperparameter tuning and model evaluation are crucial steps in the development of machine learning models. Proper tuning ensures optimal performance, while thorough evaluation helps understand the model's strengths and weaknesses.

Hyperparameter Tuning

Hyperparameters are settings that define the model structure and how it is trained, such as learning rate, batch size, number of layers, and units per layer. Unlike parameters learned during training, hyperparameters need to be set before the training process begins.

Importance of Hyperparameter Tuning

Effective hyperparameter tuning can significantly improve model performance. Poorly chosen hyperparameters can lead to underfitting or overfitting, resulting in a model that performs poorly on unseen data.

Techniques for Hyperparameter Tuning

1. Grid Search: Exhaustively searches over a specified hyperparameter grid.

2. Random Search: Samples hyperparameters randomly from a defined range.

3. Bayesian Optimization: Uses probabilistic models to find the optimal hyperparameters.

4. Hyperband: Combines random search and early stopping to efficiently find optimal hyperparameters.

Grid Search

Grid search is a brute-force technique that searches over a predefined grid of hyperparameters.

from sklearn.model_selection import GridSearchCV
from sklearn.ensemble import RandomForestClassifier

param_grid = {
    'n_estimators': [100, 200, 300],
    'max_depth': [10, 20, 30],
    'min_samples_split': [2, 5, 10]
}

grid_search = GridSearchCV(estimator=RandomForestClassifier(), param_grid=param_grid, cv=3, n_jobs=-1, verbose=2)
grid_search.fit(X_train, y_train)

print("Best Hyperparameters:", grid_search.best_params_)

Random Search

Random search samples hyperparameters from a specified distribution, which can be more efficient than grid search.

from sklearn.model_selection import RandomizedSearchCV
from sklearn.ensemble import RandomForestClassifier
from scipy.stats import randint

param_dist = {
    'n_estimators': randint(100, 500),
    'max_depth': randint(10, 50),
    'min_samples_split': randint(2, 11)
}

random_search = RandomizedSearchCV(estimator=RandomForestClassifier(), param_distributions=param_dist, n_iter=100, cv=3, n_jobs=-1, verbose=2)
random_search.fit(X_train, y_train)

print("Best Hyperparameters:", random_search.best_params_)

Bayesian Optimization

Bayesian optimization uses a surrogate model to estimate the performance of hyperparameters and efficiently searches the space.

from skopt import BayesSearchCV
from sklearn.ensemble import RandomForestClassifier

param_space = {
    'n_estimators': (100, 500),
    'max_depth': (10, 50),
    'min_samples_split': (2, 11)
}

bayes_search = BayesSearchCV(estimator=RandomForestClassifier(), search_spaces=param_space, n_iter=32, cv=3, n_jobs=-1, verbose=2)
bayes_search.fit(X_train, y_train)

print("Best Hyperparameters:", bayes_search.best_params_)

Hyperband

Hyperband combines random search with early stopping to find the best hyperparameters more efficiently.

from keras_tuner.tuners import Hyperband
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

def build_model(hp):
    model = Sequential()
    model.add(Dense(units=hp.Int('units', min_value=32, max_value=512, step=32), activation='relu', input_shape=(input_dim,)))
    model.add(Dense(1, activation='sigmoid'))
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
    return model

tuner = Hyperband(build_model, objective='val_accuracy', max_epochs=10, factor=3, directory='my_dir', project_name='helloworld')
tuner.search(X_train, y_train, epochs=50, validation_split=0.2)

print("Best Hyperparameters:", tuner.get_best_hyperparameters()[0].values)

Model Evaluation

Model evaluation involves assessing the performance of a trained model using various metrics. This helps determine how well the model generalizes to new, unseen data.

Evaluation Metrics

1. Accuracy: Proportion of correctly predicted instances.

2. Precision: Proportion of true positives among the predicted positives.

3. Recall: Proportion of true positives among the actual positives.

4. F1 Score: Harmonic mean of precision and recall.

5. ROC-AUC: Area under the Receiver Operating Characteristic curve, measuring the trade-off between true positive rate and false positive rate.

6. Mean Squared Error (MSE): Average of the squared differences between predicted and actual values (for regression).

7. Mean Absolute Error (MAE): Average of the absolute differences between predicted and actual values (for regression).

Cross-Validation

Cross-validation is a technique for assessing model performance by splitting the data into multiple folds and training/testing the model on these folds. Common methods include k-fold cross-validation and stratified k-fold cross-validation.

from sklearn.model_selection import cross_val_score
from sklearn.ensemble import RandomForestClassifier

model = RandomForestClassifier(n_estimators=100)
scores = cross_val_score(model, X, y, cv=5, scoring='accuracy')

print("Cross-Validation Scores:", scores)
print("Mean Accuracy:", scores.mean())

Confusion Matrix

A confusion matrix provides a detailed breakdown of model predictions, showing the counts of true positives, true negatives, false positives, and false negatives.

from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

y_pred = model.predict(X_test)
cm = confusion_matrix(y_test, y_pred)
disp = ConfusionMatrixDisplay(confusion_matrix=cm)
disp.plot()
plt.show()

ROC Curve and AUC

The ROC curve plots the true positive rate against the false positive rate at various threshold settings. The AUC represents the area under this curve.

from sklearn.metrics import roc_curve, roc_auc_score
import matplotlib.pyplot as plt

y_pred_proba = model.predict_proba(X_test)[:, 1]
fpr, tpr, _ = roc_curve(y_test, y_pred_proba)
auc = roc_auc_score(y_test, y_pred_proba)

plt.plot(fpr, tpr, label=f'ROC Curve (AUC = {auc:.2f})')
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('Receiver Operating Characteristic')
plt.legend(loc='lower right')
plt.show()

Conclusion - Part II

Implementing advanced model architectures with TensorFlow encompasses a broad range of techniques and methodologies, each crucial for developing robust, efficient, and high-performing machine learning models. From setting up the development environment to fine-tuning hyperparameters and evaluating models, every step plays a vital role in the model development lifecycle.

Key Takeaways

1. Implementing Attention Mechanisms: Attention mechanisms, especially in the context of the Transformer architecture, have revolutionized the way models handle sequential data. By enabling models to focus on relevant parts of the input, attention mechanisms significantly enhance the capability of models to understand complex dependencies.

2. Building Generative Models: Generative models like Variational Autoencoders (VAEs) and Generative Adversarial Networks (GANs) open up new possibilities in data generation and augmentation. These models are particularly powerful in applications such as image synthesis, data augmentation, and creative AI tasks.

3. Hyperparameter Tuning and Model Evaluation: Hyperparameter tuning is a critical step in optimizing model performance. Techniques like grid search, random search, Bayesian optimization, and Hyperband provide systematic approaches to finding the best hyperparameters. Model evaluation metrics and methods ensure that the models are not only accurate but also generalize well to unseen data.

Final Thoughts

Building and deploying advanced model architectures with TensorFlow requires a blend of theoretical knowledge and practical skills. By understanding and applying the concepts covered in this tutorial, developers can build sophisticated models capable of solving a wide range of real-world problems. The journey from setting up the environment to fine-tuning hyperparameters and evaluating model performance is iterative and requires continuous learning and experimentation. With TensorFlow’s powerful capabilities and a systematic approach, the possibilities for innovation in machine learning are vast and exciting.

Embarking on this journey will not only enhance your technical skills but also enable you to contribute to the rapidly advancing field of artificial intelligence, pushing the boundaries of what is possible with machine learning.

Implementing advanced model architecture with TensorFlow is a crucial aspect of building powerful and effective machine learning models. TensorFlow, an open-source machine learning library, provides a versatile framework for designing, training, and deploying various neural network architectures.

In the rapidly evolving field of machine learning, staying ahead often requires a deep understanding of advanced model architectures. In this article, we will take a dive through creating sophisticated models using TensorFlow, exploring foundational concepts and advanced techniques.

Importance of Advanced Model Architecture

While basic models are suitable for simple tasks, advanced model architectures are essential for tackling more complex problems and achieving state-of-the-art performance. As machine learning tasks become increasingly sophisticated, the need for specialized architectures, such as neural networks, attention mechanisms, generative models, hyper-parameter tuning and model evaluation.

Setting Up Your Environment

Before diving into advanced model architectures, we need to set up our environment. We'll use TensorFlow, a powerful open-source library for machine learning.

We'll use Jupyter notebooks for this article. If you don't have Jupyter installed, you can install it using pip:

pip install jupyter

Install TensorFlow

TensorFlow can be installed in Python easily, just like any other module, with a terminal command using pip, the package manager for Python. Open a terminal or command prompt and enter the following command:

pip install tensorflow

Note: This general installation command may not work for all operating systems.

macOS

To install TensorFlow that is optimized for Apple's macOS processors (especially M1 and M2 chips) without going through the troubles of using the general installation, the following command is used:

pip install tensorflow-macos

Windows & Ubuntu without GPU

You can install the CPU version of TensorFlow on Windows & Ubuntu if you do not have an external GPU installed or wish to use the CPU:

pip install tensorflow-cpu

Note: Training neural networks with CPUs has performance issues in comparison to powerful GPUs.

TensorFlow also provides a high-level API, called Keras, which can simplify the creation and training of machine learning models. You can install Keras using the following command:

pip install tensorflow-keras

To start a Jupyter notebook, run:

jupyter notebook

Importing Necessary Libraries

In your Jupyter notebook, start by importing the necessary libraries:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import matplotlib.pyplot as plt

Understanding Custom Layers and Models

Custom layers and models in TensorFlow give you the flexibility to build complex and tailored machine learning models. This section explores the creation of custom layers, custom models, and demonstrates how they can be integrated into a neural network.

Creating Custom Layers

Custom layers allow you to encapsulate custom operations in a reusable and modular way. TensorFlow's Layer class is the building block for creating custom layers. Let's break down the process of creating a custom layer.

Basic Structure of a Custom Layer

A custom layer typically involves defining the following components:

1. Initialization (__init__ method): This method is where you define the attributes of the layer.

2. Build (build method): This method is where you define the weights and other variables that the layer will use.

3. Call (call method): This method contains the forward pass logic, specifying how the layer should process its inputs to produce outputs.

Custom Dense Layer

Here's an example of a custom dense (fully connected) layer that adds a scalar value to the input.

class MyCustomLayer(layers.Layer):
    def __init__(self, units=32, activation=None):
        super(MyCustomLayer, self).__init__()
        self.units = units
        self.activation = keras.activations.get(activation)

    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                                 initializer='random_normal',
                                 trainable=True)
        self.b = self.add_weight(shape=(self.units,),
                                 initializer='zeros',
                                 trainable=True)

    def call(self, inputs):
        return self.activation(tf.matmul(inputs, self.w) + self.b)

inputs = keras.Input(shape=(784,))
x = MyCustomLayer(64, activation='relu')(inputs)
outputs = MyCustomLayer(10, activation='softmax')(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

• Initialization (__init__ method): Initializes the number of units and activation function for the layer.

• Build (build method): Defines the weights (self.w) and biases (self.b) for the layer.

• Call (call method): Implements the forward pass, applying the weights, biases, and activation function to the inputs.

Creating Custom Models

Custom models allow you to define complex architectures beyond the sequential and functional APIs. TensorFlow's Model class is used to create custom models by subclassing it and defining the forward pass logic in the call method.

Custom Model with Functional API

Let's create a custom model using the functional API, incorporating our custom layer.

class MyCustomModel(keras.Model):
    def __init__(self, units=32, num_classes=10):
        super(MyCustomModel, self).__init__()
        self.dense1 = MyCustomLayer(units, activation='relu')
        self.dense2 = MyCustomLayer(num_classes, activation='softmax')

    def call(self, inputs):
        x = self.dense1(inputs)
        return self.dense2(x)

model = MyCustomModel(units=64, num_classes=10)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

• Initialization (__init__ method): Initializes two instances of the custom layer (dense1 and dense2) with specified units and activation functions.

• Call (call method): Implements the forward pass, applying the first custom layer followed by the second.

Custom Training Loop

For more control over the training process, you can write a custom training loop. This allows you to customize every aspect of the training process, including the forward pass, backward pass, and optimization.

Custom Training Loop

Here's an example of a custom training loop for our custom model.

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train, x_test = x_train.astype("float32") / 255.0, x_test.astype("float32") / 255.0
x_train = x_train.reshape(-1, 784)
x_test = x_test.reshape(-1, 784)

epochs = 5
batch_size = 64
optimizer = keras.optimizers.Adam()
loss_fn = keras.losses.SparseCategoricalCrossentropy()

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)

model = MyCustomModel(units=64, num_classes=10)

for epoch in range(epochs):
    print(f"Epoch {epoch+1}/{epochs}")
    for step, (x_batch, y_batch) in enumerate(train_dataset):
        with tf.GradientTape() as tape:
            logits = model(x_batch, training=True)
            loss = loss_fn(y_batch, logits)
        gradients = tape.gradient(loss, model.trainable_weights)
        optimizer.apply_gradients(zip(gradients, model.trainable_weights))

        if step % 100 == 0:
            print(f"Step {step}, Loss: {loss.numpy():.4f}")

• Data Preparation: Loads and preprocesses the MNIST dataset.

• Training Loop: Iterates over epochs and batches, computes gradients, and updates weights using the optimizer.

Implementing Neural Network Architectures

In this section, we'll dive into various neural network architectures that are fundamental in deep learning: Convolutional Neural Networks (CNNs), Recurrent Neural Networks (RNNs), and using Transfer Learning with Pre-trained Models. Each architecture is tailored for specific types of tasks and data, and we'll explore how to implement them in TensorFlow.

Convolutional Neural Networks (CNNs)

Convolutional Neural Networks (CNNs) are specialized for processing grid-like data, such as images. They are designed to automatically and adaptively learn spatial hierarchies of features through backpropagation. CNNs are composed of convolutional layers, pooling layers, and fully connected layers.

Key Components of CNNs

1. Convolutional Layers: Apply convolutional operations to the input, using filters to extract features like edges, textures, and patterns.

2. Pooling Layers: Reduce the spatial dimensions (width and height) of the data, typically using max pooling or average pooling.

3. Fully Connected Layers: Connect every neuron in one layer to every neuron in the next layer, used for classification tasks.

Implementing a Simple CNN for Image Classification

Let's implement a simple CNN for classifying images from the CIFAR-10 dataset.

import tensorflow as tf
from tensorflow.keras import layers, models

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = models.Sequential()
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))

model.add(layers.Flatten())
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(10, activation='softmax'))

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

model.fit(x_train, y_train, epochs=10, validation_data=(x_test, y_test))

• Convolutional Layers: The first layer has 32 filters of size 3x3 and ReLU activation. This is followed by a max pooling layer. This pattern repeats, with increasing filter sizes.

• Fully Connected Layers: After flattening the output from the convolutional layers, we add a dense layer with 64 units and ReLU activation, followed by a dense layer with 10 units and softmax activation for classification.

Recurrent Neural Networks (RNNs)

Recurrent Neural Networks (RNNs) are designed for sequential data, such as time series or text. They maintain a hidden state that captures information about previous inputs, making them suitable for tasks where context or order matters.

Key Components of RNNs

1. Recurrent Layers: Process each element of the sequence, maintaining a hidden state that is updated at each step.

2. LSTM and GRU: Variants of RNNs that use gating mechanisms to better capture long-range dependencies and mitigate the vanishing gradient problem.

Implementing an RNN for Text Classification

Let's implement an RNN for classifying sentiments in text data using the IMDB dataset.

import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.datasets import imdb
from tensorflow.keras.preprocessing import sequence

max_features = 10000
maxlen = 500
(x_train, y_train), (x_test, y_test) = imdb.load_data(num_words=max_features)
x_train = sequence.pad_sequences(x_train, maxlen=maxlen)
x_test = sequence.pad_sequences(x_test, maxlen=maxlen)

model = models.Sequential()
model.add(layers.Embedding(max_features, 128))
model.add(layers.SimpleRNN(128))
model.add(layers.Dense(1, activation='sigmoid'))

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

model.fit(x_train, y_train, epochs=10, validation_data=(x_test, y_test))

• Embedding Layer: Converts the input sequences into dense vectors of fixed size.

• SimpleRNN Layer: Processes the sequence data, maintaining a hidden state that captures information about the sequence.

• Dense Layer: Outputs a single value with sigmoid activation for binary classification (positive or negative sentiment).

Transfer Learning with Pre-trained Models

Transfer learning leverages pre-trained models, usually trained on large datasets, and fine-tunes them for a specific task. This approach is beneficial when you have limited data.

Steps in Transfer Learning

1. Select a Pre-trained Model: Choose a model pre-trained on a large dataset, such as ImageNet.

2. Load the Pre-trained Model: Load the model with pre-trained weights, excluding the top layers.

3. Add Custom Layers: Add new layers for your specific task.

4. Freeze the Base Layers: Freeze the weights of the pre-trained layers.

5. Compile and Train the Model: Compile and train the model on your dataset.

6. Unfreeze and Fine-tune: Optionally, unfreeze some layers and fine-tune the entire model.

Transfer Learning with ResNet50

Let's implement transfer learning using the ResNet50 model for image classification on the CIFAR-10 dataset.

from tensorflow.keras.applications import ResNet50
from tensorflow.keras import layers, models
from tensorflow.keras.optimizers import Adam

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(32, 32, 3))
x = base_model.output
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dense(1024, activation='relu')(x)
x = layers.Dense(10, activation='softmax')(x)

model = models.Model(inputs=base_model.input, outputs=x)

for layer in base_model.layers:
    layer.trainable = False

model.compile(optimizer=Adam(), loss='sparse_categorical_crossentropy', metrics=['accuracy'])

model.fit(x_train, y_train, epochs=10, validation_data=(x_test, y_test))

for layer in base_model.layers[-10:]:
    layer.trainable = True

model.compile(optimizer=Adam(learning_rate=1e-5), loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(x_train, y_train, epochs=10, validation_data=(x_test, y_test))

• Load the Pre-trained Model: The ResNet50 model is loaded with weights pre-trained on ImageNet, excluding the top layers.

• Add Custom Layers: Global average pooling is added to reduce the spatial dimensions, followed by dense layers for classification.

• Freeze the Base Layers: The base layers of ResNet50 are frozen to retain the learned features.

• Compile and Train: The model is compiled and initially trained. Then, some layers are unfrozen for fine-tuning with a lower learning rate.

Conclusion - Part I

Implementing advanced model architectures with TensorFlow is a multifaceted process that requires a solid understanding of various components and techniques. In this first part, we have laid the groundwork for developing sophisticated machine learning models by covering the following key areas:

Key Takeaways

1. Setting Up Your Environment: Establishing a robust and efficient development environment is the first step towards successful model implementation. Essential tools such as TensorFlow, Keras, and Jupyter Notebooks provide a strong foundation for experimentation and development.

2. Custom Layers and Models: Creating custom layers and models allows developers to tailor neural network architectures to specific tasks and data types. This customization enhances the flexibility and effectiveness of the models, enabling them to tackle complex challenges more efficiently.

3. Implementing Neural Network Architectures: Understanding and implementing various neural network architectures is crucial for addressing different types of data and tasks. Convolutional Neural Networks (CNNs) excel at image processing, while Recurrent Neural Networks (RNNs) are well-suited for sequential data. Each architecture has its strengths and applications, and mastering them is essential for building effective models.

4. Transfer Learning with Pre-trained Models: Transfer learning leverages existing, well-trained models to accelerate development and improve performance. By fine-tuning pre-trained models on new datasets, developers can achieve high accuracy and efficiency with less training time and data. This approach is particularly beneficial when dealing with limited data or complex tasks.

Let's wait for Part II soon, shall we?

Text classification is a basic job in natural language processing (NLP) that is used in sentiment analysis, spam detection, and content categorization. Transformer-based models, like BERT (Bi-directional Encoder Representations from Transformers), have become popular recently because of their outstanding performance in different NLP tasks.

In this article, we'll explore how to implement text classification using BERT and the KerasNLP library, providing examples and code snippets to guide you through the process.

Understanding BERT

BERT, introduced by Google in 2018, is a pre-trained transformer-based model created for understanding natural language. Unlike traditional models that analyze text in one direction, BERT looks at context from both sides, which helps it capture complex relationships within sentences effectively.

BERT Architecture

BERT's architecture consists of layers of attention mechanisms and feedforward neural networks. It employs a transformer encoder stack, allowing it to learn contextualized representations of words. The model is pre-trained on large corpora, gaining a deep understanding of language nuances.

Tokenization with BERT

Before delving into text classification, it's crucial to understand tokenization, a process that breaks down text into smaller units, such as words or subwords. BERT utilizes WordPiece tokenization, which divides text into subword tokens, enhancing its ability to handle out-of-vocabulary words.

Setting Up the Environment

To get started with BERT-based text classification, you need to set up your Python environment. Ensure you have the required libraries installed:

pip install tensorflow
pip install keras-nlp
pip install transformers

These packages include TensorFlow, KerasNLP, and the Hugging Face Transformers library, which provides pre-trained BERT models.

Import the libraries

from keras_nlp import load_bert_model
import pandas as pd
from sklearn.model_selection import train_test_split

from keras.layers import Input, Dense
from keras.models import Model
from keras.optimizers import Adam

from keras_nlp import Tokenizer

from keras_nlp import load_bert_finetuned_model

Loading BERT Model with KerasNLP

KerasNLP simplifies the process of working with BERT models in Keras. Let's load a pre-trained BERT model using KerasNLP:

model_name = 'bert-base-uncased'
bert_model = load_bert_model(model_name)

Note: You can choose other variants based on your requirements, such as multilingual models or models fine-tuned for specific tasks.

Text Classification

Now, let's move on to text classification using BERT. For this example, we'll create a binary sentiment analysis model. Assume you have a dataset with labeled sentiments (positive or negative). First, load and preprocess the data

data = pd.read_csv('sentiment_data.csv')

train_data, test_data = train_test_split(data, test_size=0.2, random_state=42)

tokenizer = Tokenizer(model_name)
X_train = tokenizer.tokenize(train_data['text'].tolist())
X_test = tokenizer.tokenize(test_data['text'].tolist())

y_train = train_data['sentiment'].map({'negative': 0, 'positive': 1}).values
y_test = test_data['sentiment'].map({'negative': 0, 'positive': 1}).values

In this example, we assume that your dataset has a 'text' column containing the text data and a 'sentiment' column with labels ('negative' or 'positive'). Adjust the column names based on your dataset structure.

Building the BERT Text Classification Model

Now, let's build the BERT-based text classification model using Keras:

input_layer = Input(shape=(tokenizer.max_seq_length,), dtype='int32')

bert_output = bert_model(input_layer)

output_layer = Dense(1, activation='sigmoid')(bert_output['pooled_output'])

model = Model(inputs=input_layer, outputs=output_layer)

model.compile(optimizer=Adam(learning_rate=2e-5), loss='binary_crossentropy', metrics=['accuracy'])

This code snippet creates a simple neural network for sentiment analysis. The BERT output is fed into a dense layer with a sigmoid activation function for binary classification. Adjust the architecture based on your specific task and requirements.

Training the BERT Text Classification Model

Now, let's train the BERT text classification model using the prepared data:

model.fit(X_train, y_train, epochs=3, batch_size=32, validation_split=0.1)

loss, accuracy = model.evaluate(X_test, y_test)
print(f'Test Loss: {loss}, Test Accuracy: {accuracy}')

This code snippet trains the model for three epochs with a batch size of 32 and validates on a 10% subset of the training data. After training, it evaluates the model on the test set, providing insights into its performance.

Fine-Tuning BERT for Specific Tasks

While the above example demonstrates a basic BERT text classification model, fine-tuning allows you to adapt BERT to specific tasks or domains. Fine-tuning involves training the pre-trained BERT model on a task-specific dataset, enabling it to learn task-specific features.

Loading a Fine-Tuned BERT Model

Assuming you have a fine-tuned BERT model saved, you can load it using KerasNLP:

fine_tuned_model_path = 'path/to/fine_tuned_model'
fine_tuned_model = load_bert_finetuned_model(fine_tuned_model_path)

Replace 'path/to/fine_tuned_model' with the actual path to your fine-tuned BERT model.

Fine-Tuning BERT for Text Classification

Let's explore how to fine-tune BERT for text classification using KerasNLP. Assume you have a task-specific dataset with text and corresponding labels:

task_data = pd.read_csv('task_specific_data.csv')

X_task = tokenizer.tokenize(task_data['text'].tolist())

y_task = task_data['label'].values

Now, fine-tune the BERT model on your task-specific dataset:

fine_tuned_model.fit(X_task, y_task, epochs=5, batch_size=16, validation_split=0.1)

fine_tuned_model.save('path/to/save/fine_tuned_model')

This code snippet fine-tunes the BERT model on the task-specific dataset for five epochs with a batch size of 16, validating on a 10% subset. After fine-tuning, it saves the model for future use.

Conclusion

In this article, we explored text classification using BERT and KerasNLP. We explained the fundamentals of BERT, prepared the environment, loaded a pre-trained BERT model, and created a basic text classification model. Furthermore, we talked about fine-tuning BERT for particular tasks, offering code snippets and examples to help you along the way.

Implementing text classification with BERT offers numerous opportunities for NLP applications. Whether you're focusing on sentiment analysis, spam detection, or any classification task, using BERT can greatly improve the accuracy and reliability of your models. As NLP progresses, keeping abreast of the newest developments and integrating them into your projects will keep you ahead in this dynamic and thrilling field.

TensorFlow is a free and open-source software library that can be used to build machine learning models. It includes the Keras API, which provides a user-friendly interface for building models. Machine learning engineers make decisions about the architecture of a model based on the type of data they are working with, the task they are trying to accomplish, and the resources they have available. The best way to learn how to build models in TensorFlow is to start with a simple task and then gradually work your way up to more complex tasks.

Why and How?

Machine learning engineers consider the type of problem, the properties of the data, and the intended performance of the model when choosing a model architecture. Here are some tips to help you understand the decisions they make:

Start with a simple architecture: It is often a good idea to start with a simple architecture when building a new model and add complexity as needed. This allows you to quickly test and improve your ideas.

Experiment with different architectures: There is no one-size-fits-all answer to model architecture. It is important to try different architectures to see which best solves your problem.

Use prior knowledge: If you know about the problem you are trying to solve or the data you are using, you can use this information to inform your model architecture decisions.

Stay up-to-date with the field: Machine learning is a constantly evolving field, with new methods and architectures being developed all the time. Staying up-to-date with the latest research can help you make informed decisions about your model design.

Building the model

Let's say you are building a model to classify images of cats and dogs. You could start with a simple architecture, such as a convolutional neural network (CNN). You could then experiment with different architectures, such as a recurrent neural network (RNN) or a long short-term memory (LSTM) network. You could also use prior knowledge about the problem, such as the fact that cats and dogs have different fur patterns, to inform your architectural decisions. Finally, you could stay up-to-date with the latest research on image classification to find new and improved architectures.

Import libraries

# import libraries already installed 
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

Prepare your data

Load your dataset using appropriate methods (e.g., tf.keras.datasets, pandas, etc.).

Preprocess your data if needed (e.g., normalization, scaling, feature engineering).

Split your data into training and testing sets.

Model architecture

# Create a sequential model
model = Sequential()

# Add layers to the model
model.add(Flatten(input_shape=(28, 28)))  
model.add(Dense(128, activation='relu'))                      
model.add(Dense(10))

We are creating a basic neural network model using TensorFlow's Keras API. The first layer is a Flatten layer that takes an input image with a shape of (28, 28) and flattens it into a 1D array. The subsequent layer is a Dense layer with 128 neurons, using the ReLU activation function. Finally, we have another Dense layer with 10 neurons.

Our choice of the number of layers and neurons was based on previous knowledge and experimentation. For example, 128 neurons in the hidden layer have been shown to perform well on similar problems in the past. Similarly, using the ReLU activation function is a common choice as it has been proven to be effective in practice. For more detailed explanations and information, please refer to the provided link.

After constructing the model, the next step is to compile it. This involves instructing TensorFlow on how we want it to learn. Our ultimate goal is to enable our program to learn and become more intelligent to tackle future challenges.

Compiling a Model in TensorFlow

Building a TensorFlow model is like building a Jenga tower, where the different layers of the model correspond to the different types of blocks in the tower. When building the model, we check that the tower is stable and that all the bricks are neatly stacked. Compiling the model is like adding the finishing touches to the tower.

To compile a machine learning model, we select components like the loss function (how well the model is performing) and the optimizer (which helps adjust the blocks to make the tower more stable). The loss function acts as a benchmark, while the optimizer acts as a shovel. Machine learning engineers select the best-performing loss function and optimizer for their problem.

Here’s an example of compiling a model in TensorFlow:

# Configure learning process
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

In TensorFlow, we compile a model to set up the loss function, optimizer, and metrics. This is like ensuring that all the Jenga blocks are properly placed. After creating the model, we can fit it with data to train it.

In this example, we are telling TensorFlow that we want to use categorical_crossentropy as our loss function and adam as our optimizer. We are also saying that we want to keep track of how accurate our model is by including accuracy in our list of metrics.

Fitting a Model in TensorFlow

When working with TensorFlow, training a model is similar to playing Jenga. You must maintain the balance of your tower of blocks, with each layer representing a different level of the tower. As you add or remove blocks, you assess how well the tower can remain stable.

In the field of machine learning, fitting a model involves providing it with data to learn from. The model examines the data and tries to make predictions, then measures the accuracy of those predictions. If the results are not up to standard, the model tweaks its settings and tries again, aiming to enhance its accuracy with each attempt.

# Train the model on your training data
model.fit(x_train, y_train, epochs=10, batch_size=32)

Training a TensorFlow model involves feeding it data (x_train, y_train) and letting it learn through multiple passes (epochs). We test its performance on unseen data (x_test, y_test) to gauge its progress.

Why is fitting important?

Model fitting is the core of machine learning. Just like a poorly stacked Jenga tower, a poorly fitted model won't be reliable for real-world decisions. Fitting finds the best internal settings (hyperparameters) for your data, allowing the model to extract key information and make accurate predictions.

Think of it as automated tuning

Fitting automatically adjusts your model's parameters to optimally solve your specific problem. This ensures high accuracy and eliminates manual parameter tweaking.

Evaluate your model

# Evaluate model performance on test data
test_loss, test_acc = model.evaluate(x_test, y_test)
print('Test accuracy:', test_acc)

Why Choose TensorFlow?

Flexibility and Versatility

TensorFlow supports various deep learning tasks like image recognition, natural language processing, and time series forecasting. It caters to a wide range of applications, making it a versatile choice for different projects. Its diverse backend options, including Python, C++, and Java, allow for integration with various existing systems and tools, enhancing flexibility.

Scalability and Performance

TensorFlow can handle large datasets and complex models efficiently, thanks to its distributed computing capabilities. This allows scaling up your training process for faster model development and deployment. Its integration with various cloud platforms like Google Cloud TPUs and NVIDIA GPUs further boosts performance and scalability.

Eager Execution and Debugging

TensorFlow offers eager execution, enabling line-by-line code evaluation and debugging. This makes it easier to understand and troubleshoot your model's behavior, leading to faster development cycles. Visualization tools like TensorBoard provide insights into your model's training process, allowing you to monitor performance and identify potential issues.

Continuous Development and Innovation

TensorFlow is constantly evolving, with regular updates and new features. This ensures access to cutting-edge advancements in the field of deep learning and machine learning. The active development team and community contribute to ongoing improvements in stability, performance, and usability, making TensorFlow a reliable and future-proof choice.

Disadvantages

Steep Learning Curve

TensorFlow can have a steeper learning curve compared to some other frameworks, especially for beginners. It's complex API and diverse functionalities require a dedicated effort to master. While the extensive community and resources can help, initial setup and configuration might require additional time and effort.

Resource Intensity

Training complex models in TensorFlow can be resource-intensive, demanding powerful hardware and computing resources. This can be a constraint for smaller projects or those with limited budgets. Cloud platforms can alleviate this issue, but their costs need to be factored in when making the decision.

Debugging Challenges

Debugging complex models in TensorFlow can be challenging due to its intricate architecture and data flow. While eager execution helps, identifying the root cause of issues might require advanced knowledge and expertise. Investing in proper monitoring and logging practices can help mitigate this challenge.

Potential for Overfitting

TensorFlow's flexibility allows for building powerful models, but it also increases the risk of overfitting. This occurs when the model memorizes the training data instead of learning generalizable patterns. Techniques like regularization and early stopping can help prevent overfitting, but careful tuning might be necessary.

Conclusion

You've taken a major step into the world of deep learning by understanding the fundamentals of building, compiling, and fitting models with TensorFlow. This journey may have its challenges, but the rewards are significant – the ability to unlock powerful insights from your data and solve complex problems.

Here are some key takeaways to keep in mind as you continue your learning journey:

Start small and iterate: Begin with simple models and gradually increase complexity as you gain confidence. Experiment with different architectures and hyperparameters to see their impact on performance.

Leverage the community: Don't hesitate to seek help from the vast TensorFlow community. Utilize online resources, forums, and documentation to troubleshoot problems and learn from others' experiences.

Practice makes perfect: The more you train models, the better you'll understand their behavior and potential pitfalls. Use diverse datasets and tasks to hone your skills and become a well-rounded machine learning practitioner.

Stay curious and engaged: The field of deep learning is constantly evolving, with new tools and techniques emerging regularly. Keep up with the latest advancements and be open to exploring new ideas to stay ahead of the curve.

Remember, building and training effective models is not just about writing code. It's about understanding the problem you're trying to solve, choosing the right tools, and iteratively refining your approach. With dedication and a curious mind, you can harness the power of TensorFlow to build impactful solutions and become a valuable asset in the world of AI and machine learning.

We won some. We lost some. There's nothing else I can add.

I wish you the very best of the 366 days ahead.

With love,

Wes.

Setting up TensorBoard in Google Colab can be incredibly useful for visualizing your machine learning model's training progress and performance. TensorBoard is a powerful tool that helps you monitor various metrics, visualize model architectures, and gain insights into your model's behavior. Here's a step-by-step tutorial with examples and code snippets to guide you through the process.

Import Necessary Libraries

First, you need to import the required libraries. Make sure you have TensorFlow installed in your Colab environment.

import tensorflow as tf
from tensorboard import notebook

Load and Prepare Your Data

For demonstration purposes, let's use a simple dataset. Replace this with your actual dataset and preprocessing steps.

# Load and preprocess your data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0  # Normalize pixel values

Build and Compile Your Model

Again, this is just a simple example. Replace it with your actual model architecture and configuration.

model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10)
])

# Compile the model
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer='adam', loss=loss_fn, metrics=['accuracy'])

Set Up TensorBoard Callback

Now, you'll create a TensorBoard callback that will save logs for visualization.

# Define the log directory
log_dir = "/content/logs"  # You can modify this path

# Create a TensorBoard callback
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)

Train Your Model

Train your model using the fit function and include the TensorBoard callback.

# Train the model
model.fit(x_train, y_train, epochs=5, callbacks=[tensorboard_callback])

Start TensorBoard in Colab

TensorBoard can be started directly within a Colab notebook using the notebook module.

# Load TensorBoard in Colab
notebook.start('--logdir ' + log_dir)

Access and Visualize TensorBoard

After running the previous cell, you'll see a link to access TensorBoard. Click on that link to open TensorBoard within your Colab environment. You can navigate through various tabs to visualize different aspects of your training process.

Stop TensorBoard

Once you're done with TensorBoard, you can stop it using the "Stop" button in the TensorBoard UI, or you can run the following code to stop the TensorBoard instance:

notebook.stop()

Conclusion

That's it! You've successfully set up TensorBoard in Google Colab to monitor and visualize your model's training progress.

Remember that this tutorial provided a basic example. Depending on your use case, you might need to adjust the code to suit your specific model architecture, dataset, and training configuration.