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

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")

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

Python

    # 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}%")

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

Python

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

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()

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%

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