Project Kickstart: Data Loading & Preprocessing

Introduction

Welcome to the very first lesson of JAX in Action: Building an Image Classifier! Congratulations on reaching this final course in our comprehensive JAX learning path — you've come incredibly far and should be proud of your dedication and progress. Let's quickly recap the amazing journey that brought us here. In our first course, we mastered JAX fundamentals, learning how NumPy-like operations work in JAX's functional programming paradigm and discovering the power of automatic differentiation with jax.grad . Our second course took us deeper into advanced JAX concepts, where we explored efficient batching with jax.vmap, flexible data structures with PyTrees, and the performance benefits of JIT compilation. The third course equipped us with essential deep learning tools, teaching us to build neural networks from scratch using pure JAX, implement training loops with gradient descent, and solve classic problems like XOR. Finally, our fourth course introduced us to the elegant world of Flax and Optax, where we learned to construct modular neural networks, leverage sophisticated optimizers like Adam, and create production-ready training pipelines. Now, in this final course, we're ready to tackle a real-world challenge: building a complete image classification system from the ground up. Throughout the course, we'll work with a real dataset, implement convolutional neural networks, and apply everything we've learned to solve practical computer vision problems. Today's first lesson focuses on establishing our project foundation and mastering the critical first step of any machine learning project: data loading and preprocessing .

Understanding MNIST and TensorFlow Datasets

The MNIST dataset serves as our gateway into image classification: think of it as the "Hello, World!" of computer vision. MNIST contains 70,000 handwritten digit images (0-9), each measuring 28\times28 pixels in grayscale. Despite its simplicity, MNIST teaches us fundamental concepts that scale to complex real-world datasets: pixel normalization, shape manipulation, and batch processing. We'll use TensorFlow Datasets (tensorflow_datasets or tfds) to access MNIST, even though we're building with JAX. This might seem counterintuitive, but tfds provides an excellent data loading ecosystem with built-in preprocessing, caching, and performance optimizations. Many JAX practitioners use this combination because TensorFlow's data pipeline integrates seamlessly with JAX arrays through NumPy compatibility. The beauty of tfds lies in its standardized interface . Whether we're working with MNIST, CIFAR-10, or ImageNet, the loading patterns remain consistent. This consistency will serve us well when we expand to more complex datasets in later lessons. The library handles downloading, caching, and version management automatically, so we can focus on the machine learning rather than data infrastructure.

Join the 1M+ learners on CodeSignal

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

Implementing the Complete Data Pipeline

Now we'll build the full data loading pipeline that handles batching, shuffling, and performance optimization:

Python

def load_mnist_dataset(batch_size: int, shuffle_buffer_size: int = 10000, seed: int = 42):
    """Loads and prepares the MNIST dataset"""
    # Load MNIST dataset
    (ds_train, ds_test), ds_info = tfds.load(
        'mnist',
        split=['train', 'test'],
        shuffle_files=True,
        as_supervised=False,
        with_info=True,
    )

    # Preprocess training data
    ds_train = ds_train.map(
        preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)
    ds_train = ds_train.cache()
    ds_train = ds_train.shuffle(shuffle_buffer_size, seed=seed)
    ds_train = ds_train.batch(batch_size)
    ds_train = ds_train.prefetch(tf.data.AUTOTUNE)

def load_mnist_dataset(batch_size: int, shuffle_buffer_size: int = 10000, seed: int = 42):
    """Loads and prepares the MNIST dataset"""
    # Load MNIST dataset
    (ds_train, ds_test), ds_info = tfds.load(
        'mnist',
        split=['train', 'test'],
        shuffle_files=True,
        as_supervised=False,
        with_info=True,
    )

    # Preprocess training data
    ds_train = ds_train.map(
        preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)
    ds_train = ds_train.cache()
    ds_train = ds_train.shuffle(shuffle_buffer_size, seed=seed)
    ds_train = ds_train.batch(batch_size)
    ds_train = ds_train.prefetch(tf.data.AUTOTUNE)

This function demonstrates production-grade data pipeline construction. The tfds.load call retrieves both training and test splits while capturing dataset metadata in ds_info. We use as_supervised=False to maintain dictionary format for our examples, giving us explicit control over the data structure.

The training data pipeline applies several performance optimizations: map with parallel processing applies our preprocessing function efficiently, cache stores preprocessed data in memory to avoid redundant computation, shuffle randomizes example order to prevent learning biases (notice that we specify a seed), batch groups examples for efficient GPU processing, and prefetch overlaps data loading with model training to minimize idle time.

Inspecting Our Data Pipeline Results

Finally, we can test our data pipeline using the src/main.py file. We'll examine what our data pipeline produces and verify everything works correctly:

Python

from data_loader import load_mnist_dataset

def main():
    BATCH_SIZE = 32
    
    # Load the MNIST dataset iterators
    train_data_iter, test_data_iter, dataset_info = load_mnist_dataset(batch_size=BATCH_SIZE)

    print(f"Dataset features: {dataset_info.features}")
    print(f"Number of training examples: {dataset_info.splits['train'].num_examples}")
    print(f"Number of test examples: {dataset_info.splits['test'].num_examples}")

    # Fetch one batch from the training data iterator
    sample_batch = next(iter(train_data_iter))
    images = sample_batch['image']
    labels = sample_batch['label']

    print(f"Images batch shape: {images.shape}, dtype: {images.dtype}")
    print(f"Labels batch shape: {labels.shape}, dtype: {labels.dtype}")
    print(f"Example label from batch: {labels[0]}")
    print(f"Min/Max pixel values in batch: {images.min():.2f} / {images.max():.2f}")

from data_loader import load_mnist_dataset

def main():
    BATCH_SIZE = 32
    
    # Load the MNIST dataset iterators
    train_data_iter, test_data_iter, dataset_info = load_mnist_dataset(batch_size=BATCH_SIZE)

    print(f"Dataset features: {dataset_info.features}")
    print(f"Number of training examples: {dataset_info.splits['train'].num_examples}")
    print(f"Number of test examples: {dataset_info.splits['test'].num_examples}")

    # Fetch one batch from the training data iterator
    sample_batch = next(iter(train_data_iter))
    images = sample_batch['image']
    labels = sample_batch['label']

    print(f"Images batch shape: {images.shape}, dtype: {images.dtype}")
    print(f"Labels batch shape: {labels.shape}, dtype: {labels.dtype}")
    print(f"Example label from batch: {labels[0]}")
    print(f"Min/Max pixel values in batch: {images.min():.2f} / {images.max():.2f}")

This main module imports the data loading functionality and initializes the pipeline with a batch size of 32. This batch size balances memory usage with training efficiency — small enough to fit comfortably in memory while large enough to provide stable gradient estimates during training. Then, we extract a sample image and related label for inspection.

When we execute this code, we'll see comprehensive information about our data pipeline:

text

Dataset features: FeaturesDict({
    'image': Image(shape=(28, 28, 1), dtype=uint8),
    'label': ClassLabel(shape=(), dtype=int64, num_classes=10),
})
Number of training examples: 60000
Number of test examples: 10000
Images batch shape: (32, 28, 28, 1), dtype: float32
Labels batch shape: (32,), dtype: int32
Example label from batch: 5
Min/Max pixel values in batch: 0.00 / 1.00

Dataset features: FeaturesDict({
    'image': Image(shape=(28, 28, 1), dtype=uint8),
    'label': ClassLabel(shape=(), dtype=int64, num_classes=10),
})
Number of training examples: 60000
Number of test examples: 10000
Images batch shape: (32, 28, 28, 1), dtype: float32
Labels batch shape: (32,), dtype: int32
Example label from batch: 5
Min/Max pixel values in batch: 0.00 / 1.00

This output confirms our pipeline works perfectly! The dataset contains 60,000 training examples and 10,000 test examples. Each image is properly shaped as (28, 28, 1), and our preprocessing successfully normalized pixel values to the [0,1] range. The batch dimensions show we're getting 32 images per batch, exactly as configured, with the correct data types for JAX consumption.

Introduction

Understanding MNIST and TensorFlow Datasets

Join the 1M+ learners on CodeSignal

Building Our Data Loading Foundation

Implementing the Complete Data Pipeline

Completing the Data Pipeline

Inspecting Our Data Pipeline Results

Conclusion and Next Steps