Full Pipeline: Training and Evaluating the CNN

Introduction

Welcome back to the fourth lesson of JAX in Action: Building an Image Classifier! We've constructed a solid foundation over our previous lessons: efficient data loading with preprocessing pipelines, a sophisticated CNN architecture using Flax, and the essential training and evaluation utilities that form the computational core of our learning system. Now comes the exciting culmination, where we orchestrate all these components into a complete, functioning training pipeline. In this lesson, we'll build the main orchestration layer that brings everything together in a seamless training and evaluation workflow. We'll implement the full pipeline in our main module, complete with hyperparameter configuration, systematic epoch-based training loops, and comprehensive evaluation cycles that track our model's learning progress. This represents the final piece of our machine learning puzzle, transforming our modular components into a cohesive system capable of training a CNN from random initialization to high accuracy on MNIST digit classification.

Setting Up Hyperparameters

Let's begin our implementation by establishing the fundamental configuration that will support our entire training process: Pythonimport jax import jax.numpy as jnp import optax import time from data_loader import load_mnist_dataset from model import CNN from train_utils import train_step, eval_step def main(): # Hyperparameters NUM_EPOCHS = 5 BATCH_SIZE = 64 LEARNING_RATE = 1e-3 PRNG_SEED = 0 print(f"Hyperparameters: Epochs={NUM_EPOCHS}, Batch Size={BATCH_SIZE}, LR={LEARNING_RATE}") This configuration section establishes the hyperparameter foundation for our training process. We set a modest 5 epochs since we're limiting our dataset size to keep execution time reasonable for demonstration purposes. The batch size of 64 provides a good balance between computational efficiency and gradient noise, while the learning rate of 1e-3 represents a proven starting point for Adam optimization. The PRNG seed ensures reproducible results across different runs, which is crucial for debugging and comparison purposes.

Setting up Data Loading

Let's now focus on the data infrastructure: Python # Load datasets max_train_samples = 1000 max_test_samples = 200 train_data_iter, test_data_iter, dataset_info = load_mnist_dataset( batch_size=BATCH_SIZE, max_train_samples=max_train_samples, max_test_samples=max_test_samples ) num_train_batches = max_train_samples // BATCH_SIZE num_test_batches = max_test_samples // BATCH_SIZE print(f"Loaded MNIST: {num_train_batches} training batches, {num_test_batches} test batches") The data loading configuration deliberately limits our dataset size to 1,000 training samples and 200 test samples. This limitation serves a practical purpose: keeping our training demonstration fast and focused while still providing meaningful learning dynamics. We calculate the exact number of batches upfront, which will be essential for our training loop logic and ensures we know precisely how many iterations to perform in each epoch.

Join the 1M+ learners on CodeSignal

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

Initializing Model and Optimizer Components

With our data pipeline established, we now initialize the core learning components that will drive our training process:

Python

    # Initialize model, parameters, and optimizer
    prng_key = jax.random.key(PRNG_SEED)
    key_init_model, _ = jax.random.split(prng_key)

    cnn_model = CNN(num_classes=10)
    dummy_input_for_init = jnp.ones((1, 28, 28, 1), dtype=jnp.float32)
    initial_variables = cnn_model.init(key_init_model, dummy_input_for_init)
    current_params = initial_variables['params']
    
    optimizer = optax.adam(learning_rate=LEARNING_RATE)
    current_opt_state = optimizer.init(current_params)
    
    print("Model, Parameters, and Optimizer initialized")

    # Initialize model, parameters, and optimizer
    prng_key = jax.random.key(PRNG_SEED)
    key_init_model, _ = jax.random.split(prng_key)

    cnn_model = CNN(num_classes=10)
    dummy_input_for_init = jnp.ones((1, 28, 28, 1), dtype=jnp.float32)
    initial_variables = cnn_model.init(key_init_model, dummy_input_for_init)
    current_params = initial_variables['params']
    
    optimizer = optax.adam(learning_rate=LEARNING_RATE)
    current_opt_state = optimizer.init(current_params)
    
    print("Model, Parameters, and Optimizer initialized")

This initialization sequence demonstrates the usual JAX practices for model setup. We create a proper JAX random key using the modern jax.random.key function, then split it to ensure clean separation between different random operations. The model initialization requires a dummy input tensor matching our expected data shape, which triggers Flax's lazy initialization mechanism to determine all parameter dimensions automatically.

The optimizer initialization creates the internal state that Adam requires to track momentum and variance estimates for each parameter. This state will evolve throughout training, accumulating the historical information that enables Adam's adaptive learning rate behavior. By storing these components in current_params and current_opt_state, we establish the mutable state that our training loop will continuously update.

Creating a Reusable Evaluation Function

Now, let's implement a reusable evaluation function that will serve both our initial baseline assessment and our epoch-by-epoch monitoring:

Python

def evaluate_model(model_apply_fn, params, test_data_iter, num_test_batches):
    """Evaluate the model on the test dataset and return metrics"""
    total_eval_loss = 0.0
    total_correct_eval = 0
    num_eval_samples = 0
    
    current_test_iter = iter(test_data_iter)
    for _ in range(num_test_batches):
        eval_batch = next(current_test_iter)
        eval_images_jax = jnp.array(eval_batch['image'])
        eval_labels_jax = jnp.array(eval_batch['label'])
        
        batch_actual_size = eval_images_jax.shape[0]

        eval_loss_batch, eval_accuracy_batch = eval_step(
            model_apply_fn, params, eval_images_jax, eval_labels_jax
        )
        total_eval_loss += eval_loss_batch * batch_actual_size
        total_correct_eval += eval_accuracy_batch * batch_actual_size
        num_eval_samples += batch_actual_size
        
    avg_eval_loss = total_eval_loss / num_eval_samples
    avg_eval_accuracy = total_correct_eval / num_eval_samples
    
    return avg_eval_loss, avg_eval_accuracy

def evaluate_model(model_apply_fn, params, test_data_iter, num_test_batches):
    """Evaluate the model on the test dataset and return metrics"""
    total_eval_loss = 0.0
    total_correct_eval = 0
    num_eval_samples = 0
    
    current_test_iter = iter(test_data_iter)
    for _ in range(num_test_batches):
        eval_batch = next(current_test_iter)
        eval_images_jax = jnp.array(eval_batch['image'])
        eval_labels_jax = jnp.array(eval_batch['label'])
        
        batch_actual_size = eval_images_jax.shape[0]

        eval_loss_batch, eval_accuracy_batch = eval_step(
            model_apply_fn, params, eval_images_jax, eval_labels_jax
        )
        total_eval_loss += eval_loss_batch * batch_actual_size
        total_correct_eval += eval_accuracy_batch * batch_actual_size
        num_eval_samples += batch_actual_size
        
    avg_eval_loss = total_eval_loss / num_eval_samples
    avg_eval_accuracy = total_correct_eval / num_eval_samples
    
    return avg_eval_loss, avg_eval_accuracy

This function encapsulates the complete evaluation workflow with proper metric aggregation. By extracting this logic into a standalone function, we eliminate code duplication and ensure consistent evaluation behavior throughout our training pipeline. The function handles weighted averaging correctly by multiplying each batch's metrics by its actual size before accumulating, then dividing by the total number of samples. This approach gracefully handles variable batch sizes and provides accurate overall metrics.

Implementing the Main Training Loop

Now we implement the core training loop that orchestrates the learning process across multiple epochs:

Python

    print("\nStarting training...")

    # Training loop
    for epoch in range(NUM_EPOCHS):
        epoch_start_time = time.time()
        
        # Training phase
        total_train_loss = 0.0
        current_train_iter = iter(train_data_iter)
        
        for batch_idx in range(num_train_batches):
            batch = next(current_train_iter)
            images_jax = jnp.array(batch['image'])
            labels_jax = jnp.array(batch['label'])
            
            current_params, current_opt_state, train_loss_batch = train_step(
                cnn_model.apply, current_params, current_opt_state, optimizer,
                images_jax, labels_jax
            )
            total_train_loss += train_loss_batch
        
        avg_train_loss = total_train_loss / num_train_batches
        epoch_end_time = time.time()
        epoch_duration = epoch_end_time - epoch_start_time

    print("\nStarting training...")

    # Training loop
    for epoch in range(NUM_EPOCHS):
        epoch_start_time = time.time()
        
        # Training phase
        total_train_loss = 0.0
        current_train_iter = iter(train_data_iter)
        
        for batch_idx in range(num_train_batches):
            batch = next(current_train_iter)
            images_jax = jnp.array(batch['image'])
            labels_jax = jnp.array(batch['label'])
            
            current_params, current_opt_state, train_loss_batch = train_step(
                cnn_model.apply, current_params, current_opt_state, optimizer,
                images_jax, labels_jax
            )
            total_train_loss += train_loss_batch
        
        avg_train_loss = total_train_loss / num_train_batches
        epoch_end_time = time.time()
        epoch_duration = epoch_end_time - epoch_start_time

The training loop structure follows a nested iteration pattern: the outer loop progresses through epochs, while the inner loop processes individual batches within each epoch. We measure epoch duration to provide insight into training performance, which becomes crucial when scaling to larger datasets.

Each batch processing cycle converts the TensorFlow dataset outputs to JAX arrays, then calls our train_step function with the current model state. Notice how we continuously update current_params and current_opt_state — this functional approach ensures that each training step builds upon the previous one's learning. The loss accumulation allows us to compute average training loss per epoch, providing a smooth metric for monitoring training progress.

Logging the Progress

Let's complete our training pipeline by adding comprehensive logging:

Python

        print(f"Epoch {epoch+1}/{NUM_EPOCHS} | "
              f"Duration: {epoch_duration:.2f}s | "
              f"Avg Train Loss: {avg_train_loss:.4f} | "
              f"Avg Eval Loss: {avg_eval_loss:.4f} | "
              f"Eval Accuracy: {avg_eval_accuracy*100:.2f}%")

    print("\nTraining complete")
    print(f"Final Evaluation Loss: {avg_eval_loss:.4f}")
    print(f"Final Evaluation Accuracy: {avg_eval_accuracy*100:.2f}%")

if __name__ == "__main__":
    main()

        print(f"Epoch {epoch+1}/{NUM_EPOCHS} | "
              f"Duration: {epoch_duration:.2f}s | "
              f"Avg Train Loss: {avg_train_loss:.4f} | "
              f"Avg Eval Loss: {avg_eval_loss:.4f} | "
              f"Eval Accuracy: {avg_eval_accuracy*100:.2f}%")

    print("\nTraining complete")
    print(f"Final Evaluation Loss: {avg_eval_loss:.4f}")
    print(f"Final Evaluation Accuracy: {avg_eval_accuracy*100:.2f}%")

if __name__ == "__main__":
    main()

When we execute this complete training pipeline, we observe the following progression:

text

Hyperparameters: Epochs=5, Batch Size=64, LR=0.001
Loaded MNIST: 15 training batches, 3 test batches
Model, Parameters, and Optimizer initialized
Initial Evaluation | Loss: 2.3398 | Accuracy: 9.38%

Starting training...
Epoch 1/5 | Duration: 2.86s | Avg Train Loss: 2.0619 | Avg Eval Loss: 1.7568 | Eval Accuracy: 45.31%
Epoch 2/5 | Duration: 0.90s | Avg Train Loss: 1.5428 | Avg Eval Loss: 1.3433 | Eval Accuracy: 52.60%
Epoch 3/5 | Duration: 0.82s | Avg Train Loss: 1.1276 | Avg Eval Loss: 0.9805 | Eval Accuracy: 64.58%
Epoch 4/5 | Duration: 0.82s | Avg Train Loss: 0.7954 | Avg Eval Loss: 0.7252 | Eval Accuracy: 77.08%
Epoch 5/5 | Duration: 0.82s | Avg Train Loss: 0.5861 | Avg Eval Loss: 0.5480 | Eval Accuracy: 81.25%

Training complete
Final Evaluation Loss: 0.5480
Final Evaluation Accuracy: 81.25%

Hyperparameters: Epochs=5, Batch Size=64, LR=0.001
Loaded MNIST: 15 training batches, 3 test batches
Model, Parameters, and Optimizer initialized
Initial Evaluation | Loss: 2.3398 | Accuracy: 9.38%

Starting training...
Epoch 1/5 | Duration: 2.86s | Avg Train Loss: 2.0619 | Avg Eval Loss: 1.7568 | Eval Accuracy: 45.31%
Epoch 2/5 | Duration: 0.90s | Avg Train Loss: 1.5428 | Avg Eval Loss: 1.3433 | Eval Accuracy: 52.60%
Epoch 3/5 | Duration: 0.82s | Avg Train Loss: 1.1276 | Avg Eval Loss: 0.9805 | Eval Accuracy: 64.58%
Epoch 4/5 | Duration: 0.82s | Avg Train Loss: 0.7954 | Avg Eval Loss: 0.7252 | Eval Accuracy: 77.08%
Epoch 5/5 | Duration: 0.82s | Avg Train Loss: 0.5861 | Avg Eval Loss: 0.5480 | Eval Accuracy: 81.25%

Training complete
Final Evaluation Loss: 0.5480
Final Evaluation Accuracy: 81.25%

This output reveals compelling learning dynamics: starting from the expected random baseline of 9.38% accuracy, our model demonstrates dramatic improvement to 45.31% after just the first epoch. The training loss decreases steadily from 2.33 to 0.58, while evaluation accuracy continues its remarkable climb to 81.25%. The decreasing epoch duration after the first epoch reflects JAX's JIT compilation reaching a steady state after the initial compilation overhead. The consistent improvement in both training and evaluation metrics indicates healthy learning without overfitting, validating our architecture and hyperparameter choices.

Introduction

Setting Up Hyperparameters

Setting up Data Loading

Join the 1M+ learners on CodeSignal

Initializing Model and Optimizer Components

Creating a Reusable Evaluation Function

Establishing the Initial Baseline

Implementing the Main Training Loop

Continuous Evaluation with Our Reusable Function

Logging the Progress

Conclusion and Next Steps