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.

Let's begin our implementation by establishing the fundamental configuration that will support our entire training process:

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

Let's now focus on the data infrastructure:

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

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

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

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

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.

Before beginning training, it's crucial to establish a performance baseline that shows our model's capabilities with random initialization:

    # Initial evaluation before training
    initial_eval_loss, initial_eval_accuracy = evaluate_model(
        cnn_model.apply, current_params, test_data_iter, num_test_batches
    )
    print(f"Initial Evaluation | Loss: {initial_eval_loss:.4f} | Accuracy: {initial_eval_accuracy*100:.2f}%")

This initial evaluation serves multiple important purposes. First, it confirms that our evaluation pipeline works correctly before we begin the training process. Second, it provides a crucial baseline that helps us understand the magnitude of improvement achieved through training. For a 10-class classification problem like MNIST, we expect random performance to hover around 10% accuracy; any significant deviation from this baseline could indicate initialization issues or implementation bugs that should be addressed before proceeding with training.

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

    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.

The evaluation phase leverages our dedicated evaluation function for clean, consistent metric computation:

        # Evaluation phase
        avg_eval_loss, avg_eval_accuracy = evaluate_model(
            cnn_model.apply, current_params, test_data_iter, num_test_batches
        )

By using our evaluate_model function, we ensure that both our initial baseline and epoch-by-epoch evaluations follow identical logic and produce comparable results. This consistency is crucial for tracking learning progress accurately and identifying potential issues like overfitting or convergence problems. The function handles all the complexity of batch processing and metric aggregation, allowing our main training loop to focus on the high-level orchestration.

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

        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:

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.

Congratulations! You've successfully implemented a complete end-to-end training pipeline that orchestrates all the components we've built throughout this course. Your implementation demonstrates professional machine learning practices: modular design with clear separation of concerns, proper state management across training steps, comprehensive metric aggregation, and systematic progress monitoring.

The remarkable journey from 9.38% to 81.25% accuracy showcases the power of modern deep learning frameworks like JAX, Flax, and Optax working in harmony. You've experienced firsthand how a randomly initialized CNN can learn to recognize handwritten digits through the elegant interplay of forward propagation, backpropagation, and parameter optimization. Now it's time to apply these concepts in practice and solidify your understanding by implementing variations and enhancements to this foundation.