Crafting a CNN with Flax: Conv and Pooling Layers

Introduction

Welcome back to JAX in Action: Building an Image Classifier! In our first lesson, we established a solid foundation by creating implementing robust data loading with the MNIST dataset using tfds . Now that we have clean, normalized image data flowing through our pipeline, we're ready to tackle the next exciting challenge: building our first convolutional neural network . Today's lesson focuses on crafting a CNN architecture using Flax's elegant Linen API. We'll discover how to combine convolutional layers, activation functions, and pooling operations to create a network specifically designed for image recognition tasks. As you may recall from our earlier courses, Flax provides a clean, functional approach to building neural networks, and CNNs showcase this elegance beautifully. By the end of this lesson, you'll have implemented a complete CNN that can process our MNIST images and produce classification logits, setting the stage for training in upcoming lessons.

Understanding Convolutional Neural Networks for Images

Before we write any code, let's build intuition about why CNNs revolutionized computer vision . Traditional dense neural networks treat images as flat vectors, losing critical spatial relationships between pixels. Imagine trying to recognize a handwritten "8" by examining individual pixel values randomly scattered without their positions: the curved loops that define an "8" would be completely lost in this chaos! Convolutional layers preserve these vital spatial relationships by applying small filters (kernels) that scan across the image systematically. Think of each filter as a pattern detector: one might specialize in finding vertical edges, another in detecting curves, and yet another in identifying corners. As these filters slide across your image, they produce feature maps that light up wherever their specific pattern appears. The magic happens when you stack multiple filters: the first layer might detect simple edges, but by combining these edge detectors, deeper layers can recognize complete shapes like loops or lines. Pooling layers act as intelligent summarizers that complement convolutions perfectly. After detecting features, pooling reduces the spatial resolution while keeping the most important information. Max pooling, for instance, looks at small regions (like 2\times2 pixels) and keeps only the strongest signal: if any part of that region detected an edge strongly, that's what matters. This makes our network robust to small shifts in the input (a "7" is still a "7" even if shifted slightly) while dramatically reducing computation. This hierarchical feature extraction mimics how our own visual system works: from detecting edges to recognizing shapes to understanding complete objects.

Join the 1M+ learners on CodeSignal

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

Building Convolutional Blocks with Conv and Pooling

Now let's implement the heart of our CNN: the convolutional blocks that will extract meaningful features from our handwritten digits. We'll craft two convolutional blocks, each combining convolution, activation, and pooling to progressively refine our understanding of the image:

Python

@nn.compact
def __call__(self, x):
    # First convolutional block
    x = nn.Conv(features=32, kernel_size=(3, 3), padding='SAME', name='conv1')(x)
    x = nn.relu(x)
    x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2), padding='SAME')

    # Second convolutional block
    x = nn.Conv(features=64, kernel_size=(3, 3), padding='SAME', name='conv2')(x)
    x = nn.relu(x)
    x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2), padding='SAME')

@nn.compact
def __call__(self, x):
    # First convolutional block
    x = nn.Conv(features=32, kernel_size=(3, 3), padding='SAME', name='conv1')(x)
    x = nn.relu(x)
    x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2), padding='SAME')

    # Second convolutional block
    x = nn.Conv(features=64, kernel_size=(3, 3), padding='SAME', name='conv2')(x)
    x = nn.relu(x)
    x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2), padding='SAME')

Let's focus on the first convolutional layer:

Uses 32 different 3×3 pixel filters, with each filter learning to detect distinct low-level patterns (like edges or corners).
The 3×3 kernel size is chosen because it is large enough to capture meaningful patterns, while being small enough to be computationally efficient.
padding='SAME' ensures output spatial dimensions match the input by adding zeros around the edges, preserving border information.
After convolution, nn.relu introduces non-linearity, enabling the network to learn complex, non-linear patterns.

After applying convolution, we have a max pooling operation:

window_shape=(2, 2) defines the pooling window size: each 2×2 pixel region is examined to find the maximum value.
strides=(2, 2) controls how far the pooling window moves between operations; in this case, we are stepping 2 pixels in each direction, which means no overlap between pooling windows.
Together, these parameters halve both the height and width of the feature maps (since we move the 2×2 window by 2 pixels each time).
Retains the strongest activations within each region, making the network robust to small shifts in the input.

The second convolutional block mirrors the first block’s structure but doubles the feature count to 64, allowing the network to combine features from the first block into more sophisticated patterns. This progressive deepening and spatial reduction, where each block increases the number of features while reducing spatial resolution, forms the backbone of effective CNN design. By trading spatial detail for increasingly abstract and powerful feature representations, the network becomes capable of recognizing complex patterns necessary for accurate image classification.

Introduction

Understanding Convolutional Neural Networks for Images

Join the 1M+ learners on CodeSignal

Introducing Flax Linen Modules and the @nn.compact Decorator

Building Convolutional Blocks with Conv and Pooling

Adding Dense Layers for Classification

Model Initialization and Testing

Verifying Our CNN Architecture

Conclusion and Next Steps