Training & Evaluation Steps: The Core Engine

Introduction

Welcome to the third lesson of JAX in Action: Building an Image Classifier! Your journey from data preprocessing to neural architecture definition has laid a strong foundation, and now we stand at a pivotal moment. With MNIST data flowing through our pipeline and a sophisticated CNN architecture ready to learn, it's time to breathe life into our system by implementing the mechanisms that will transform random parameters into learned knowledge. This lesson dives into the core engine of machine learning : the training and evaluation functions that orchestrate the learning process. We'll craft JIT-compiled functions that handle loss computation, automatic differentiation, parameter updates, and performance assessment — the essential components that separate a static model from an active learning system. JAX's functional programming paradigm and powerful compilation capabilities will enable us to build these components with both elegance and efficiency. By lesson's end, you'll have implemented a complete training utilities module that can guide our CNN from random initialization to accurate digit recognition.

Understanding Training vs Evaluation Modes

Before we implement our training engine, it's crucial to understand the fundamental distinction between training and evaluation modes in neural networks. This distinction goes beyond simple terminology: it represents two different computational pathways that serve complementary purposes in the machine learning workflow. During training mode, our system operates as an active learner, constantly adjusting its understanding based on feedback. The training process follows a specific computational flow: forward propagation generates predictions from input data, a loss function quantifies the prediction errors, automatic differentiation computes gradients showing how to reduce these errors, and an optimizer translates gradients into parameter updates. This cycle requires maintaining additional state (like momentum in the Adam optimizer) and performing expensive gradient computations, but it's the only way our model can learn from data. Evaluation mode strips away the learning machinery to focus purely on assessment. Here, we perform only forward propagation to generate predictions, then compute metrics like accuracy and loss to gauge performance. No gradients are calculated, no parameters are updated, and no optimizer state is maintained. This streamlined approach makes evaluation much faster than training, allowing us to efficiently monitor learning progress during training or assess final model quality on test data. Think of training mode as a student actively learning and adjusting their understanding, while evaluation mode is like taking a test where knowledge is assessed but not modified.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Implementing JIT-Compiled Training Steps

Now we reach the centerpiece of our training engine: the function that orchestrates a complete parameter update cycle. This implementation showcases JAX's ability to combine multiple computational patterns into a single, efficient operation:

Python

@functools.partial(jax.jit, static_argnums=(0, 3))
def train_step(model_apply_fn, params, opt_state, optimizer, x_batch, y_batch):
    """Performs a single JIT-compiled training step"""
    def compute_loss_for_grad(inner_params, current_x_batch, current_y_batch):
        logits = model_apply_fn({'params': inner_params}, current_x_batch)
        loss = cross_entropy_loss(logits, current_y_batch)
        return loss

    loss_value, grads = jax.value_and_grad(compute_loss_for_grad)(params, x_batch, y_batch)
    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    new_params = optax.apply_updates(params, updates)
    
    return new_params, new_opt_state, loss_value

@functools.partial(jax.jit, static_argnums=(0, 3))
def train_step(model_apply_fn, params, opt_state, optimizer, x_batch, y_batch):
    """Performs a single JIT-compiled training step"""
    def compute_loss_for_grad(inner_params, current_x_batch, current_y_batch):
        logits = model_apply_fn({'params': inner_params}, current_x_batch)
        loss = cross_entropy_loss(logits, current_y_batch)
        return loss

    loss_value, grads = jax.value_and_grad(compute_loss_for_grad)(params, x_batch, y_batch)
    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    new_params = optax.apply_updates(params, updates)
    
    return new_params, new_opt_state, loss_value

The JIT compilation strategy deserves special attention. By using functools.partial with static_argnums=(0, 3), we tell JAX that the model function and optimizer are static, that is, they define the computation graph structure rather than flowing through it as data. This prevents costly recompilation when only parameter values change between training steps. The nested compute_loss_for_grad function creates a clean interface for automatic differentiation: it takes parameters as its first argument and returns a scalar loss, exactly what jax.value_and_grad expects. This function simultaneously computes the loss value (for monitoring) and its gradients (for learning) in a single forward-backward pass, maximizing computational efficiency. The Optax optimizer then transforms raw gradients into sophisticated parameter updates (incorporating momentum, adaptive learning rates, or other techniques), which we apply to obtain our evolved parameters. This functional decomposition of separating gradient computation, update calculation, and parameter application provides both clarity and flexibility.

Creating Evaluation Functions for Model Assessment

While training drives learning, evaluation provides the compass that guides our assessment of the model's behaviour. Our evaluation function distills model assessment to its essential components:

Python

@functools.partial(jax.jit, static_argnums=(0,))
def eval_step(model_apply_fn, params, x_batch, y_batch):
    """Performs a single JIT-compiled evaluation step"""
    logits = model_apply_fn({'params': params}, x_batch)
    
    # Calculate evaluation loss
    eval_loss = cross_entropy_loss(logits, y_batch)
    
    # Calculate accuracy
    predicted_class = jnp.argmax(logits, axis=-1)
    eval_accuracy = jnp.mean(predicted_class == y_batch)
    
    return eval_loss, eval_accuracy

@functools.partial(jax.jit, static_argnums=(0,))
def eval_step(model_apply_fn, params, x_batch, y_batch):
    """Performs a single JIT-compiled evaluation step"""
    logits = model_apply_fn({'params': params}, x_batch)
    
    # Calculate evaluation loss
    eval_loss = cross_entropy_loss(logits, y_batch)
    
    # Calculate accuracy
    predicted_class = jnp.argmax(logits, axis=-1)
    eval_accuracy = jnp.mean(predicted_class == y_batch)
    
    return eval_loss, eval_accuracy

This streamlined function demonstrates evaluation's computational simplicity compared to training. We perform a single forward pass to obtain logits, then compute two key metrics. The loss calculation reuses our cross-entropy function, providing a consistent measure across training and evaluation. For accuracy, we extract predicted classes using jnp.argmax (finding the digit with the highest confidence) and compare them against the true labels. The element-wise equality check produces a boolean array, whose mean gives us the fraction of correct predictions — an intuitive metric that complements the more abstract loss value. Notice that we only mark the model function as static since evaluation doesn't involve optimizers, resulting in even more efficient JIT compilation.

Integrating Training Utilities in Main Pipeline

With our core functions implemented, let's integrate them into a cohesive system that brings together data, model, and optimization components:

Python

import jax
import jax.numpy as jnp
import optax
from data_loader import load_mnist_dataset
from model import CNN
from train_utils import train_step, eval_step

def main():
    BATCH_SIZE = 8
    LEARNING_RATE = 1e-3
    prng_key = jax.random.key(42)

    # Load data
    train_data_iter, test_data_iter, _ = load_mnist_dataset(batch_size=BATCH_SIZE)
    
    # Get sample batches for testing
    sample_train_batch = next(iter(train_data_iter))
    sample_test_batch = next(iter(test_data_iter))
    
    train_images_jax = jnp.array(sample_train_batch['image'])
    train_labels_jax = jnp.array(sample_train_batch['label'])
    test_images_jax = jnp.array(sample_test_batch['image'])
    test_labels_jax = jnp.array(sample_test_batch['label'])

import jax
import jax.numpy as jnp
import optax
from data_loader import load_mnist_dataset
from model import CNN
from train_utils import train_step, eval_step

def main():
    BATCH_SIZE = 8
    LEARNING_RATE = 1e-3
    prng_key = jax.random.key(42)

    # Load data
    train_data_iter, test_data_iter, _ = load_mnist_dataset(batch_size=BATCH_SIZE)
    
    # Get sample batches for testing
    sample_train_batch = next(iter(train_data_iter))
    sample_test_batch = next(iter(test_data_iter))
    
    train_images_jax = jnp.array(sample_train_batch['image'])
    train_labels_jax = jnp.array(sample_train_batch['label'])
    test_images_jax = jnp.array(sample_test_batch['image'])
    test_labels_jax = jnp.array(sample_test_batch['label'])

This integration exemplifies modular design principles in machine learning systems. By importing our training utilities from a dedicated module, we maintain a clear separation of concerns: data loading handles input processing, the model defines architecture, and training utilities manage the learning process. The hyperparameters are defined as constants at the function's beginning, making experimentation straightforward. We convert the NumPy's iterators returned by load_mnist_dataset to JAX arrays to ensure compatibility with our JAX-based training functions.

Testing Our Training Engine

Let's validate our training engine with a comprehensive test that exercises both training and evaluation paths:

Python

    # Initialize model
    cnn_model = CNN(num_classes=10)
    key_init, _ = jax.random.split(prng_key)
    initial_variables = cnn_model.init(key_init, train_images_jax)
    current_params = initial_variables['params']

    # Initialize optimizer
    optimizer = optax.adam(learning_rate=LEARNING_RATE)
    current_opt_state = optimizer.init(current_params)

    # Test train_step
    print("\nTesting train_step...")
    updated_params, updated_opt_state, train_loss_value = train_step(
        cnn_model.apply, current_params, current_opt_state, optimizer,
        train_images_jax, train_labels_jax
    )
    
    print(f"Training loss: {train_loss_value:.4f}")
    
    param_changed = not jax.tree_util.tree_all(
        jax.tree_util.tree_map(jnp.array_equal, current_params, updated_params)
    )
    print(f"Parameters changed: {param_changed}")

    # Test eval_step  
    print("\nTesting eval_step...")
    eval_loss_value, eval_accuracy_value = eval_step(
        cnn_model.apply, updated_params, test_images_jax, test_labels_jax
    )
    
    print(f"Evaluation loss: {eval_loss_value:.4f}")
    print(f"Evaluation accuracy: {eval_accuracy_value*100:.2f}%")

This test produces output that confirms our implementation's correctness:

text

Testing train_step...
Training loss: 2.3200
Parameters changed: True

Testing eval_step...
Evaluation loss: 2.1511
Evaluation accuracy: 25.00%

The parameter change verification using JAX's tree utilities confirms that training actually updates our model. Remember that the tree_map applies element-wise comparison across our entire parameter tree, while tree_all aggregates these comparisons into a single boolean. The initial loss around 2.3 aligns perfectly with theoretical expectations: for 10-class classification with random initialization, we expect $-\ln(1/10) \approx 2.3$ . Moreover, the evaluation step also executes correctly. These metrics confirm our training engine is ready to guide the model toward meaningful learning.

# Initialize model cnn_model = CNN(num_classes=10) key_init, _ = jax.random.split(prng_key) initial_variables = cnn_model.init(key_init, train_images_jax) current_params = initial_variables['params'] # Initialize optimizer optimizer = optax.adam(learning_rate=LEARNING_RATE) current_opt_state = optimizer.init(current_params) # Test train_step print("\nTesting train_step...") updated_params, updated_opt_state, train_loss_value = train_step( cnn_model.apply, current_params, current_opt_state, optimizer, train_images_jax, train_labels_jax ) print(f"Training loss: {train_loss_value:.4f}") param_changed = not jax.tree_util.tree_all( jax.tree_util.tree_map(jnp.array_equal, current_params, updated_params) ) print(f"Parameters changed: {param_changed}") # Test eval_step print("\nTesting eval_step...") eval_loss_value, eval_accuracy_value = eval_step( cnn_model.apply, updated_params, test_images_jax, test_labels_jax ) print(f"Evaluation loss: {eval_loss_value:.4f}") print(f"Evaluation accuracy: {eval_accuracy_value*100:.2f}%")

Introduction

Understanding Training vs Evaluation Modes

Join the 1M+ learners on CodeSignal

Building the Loss Function Foundation

Implementing JIT-Compiled Training Steps

Creating Evaluation Functions for Model Assessment

Integrating Training Utilities in Main Pipeline

Testing Our Training Engine

Conclusion and Next Steps