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

Python

import flax.linen as nn
import jax.numpy as jnp

class CNN(nn.Module):
    num_classes: int = 10

    @nn.compact
    def __call__(self, x):
        # ...

import flax.linen as nn
import jax.numpy as jnp

class CNN(nn.Module):
    num_classes: int = 10

    @nn.compact
    def __call__(self, x):
        # ...

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

Python

# Flatten for dense layers
x = jnp.reshape(x, (x.shape[0], -1))

# Dense layer
x = nn.Dense(features=128, name='dense1')(x)
x = nn.relu(x)

# Output layer (logits)
x = nn.Dense(features=self.num_classes, name='output_logits')(x)

# Flatten for dense layers
x = jnp.reshape(x, (x.shape[0], -1))

# Dense layer
x = nn.Dense(features=128, name='dense1')(x)
x = nn.relu(x)

# Output layer (logits)
x = nn.Dense(features=self.num_classes, name='output_logits')(x)

Python

import jax
import jax.numpy as jnp
from model import CNN

def main():
    # Test CNN model initialization
    cnn_model = CNN(num_classes=10)
    key = jax.random.key(123)
    
    # Create dummy batch for initialization
    dummy_batch_size = 1
    dummy_image_batch = jnp.ones((dummy_batch_size, 28, 28, 1), dtype=jnp.float32)
    print(f"Dummy image batch shape: {dummy_image_batch.shape}")

    # Initialize parameters
    variables = cnn_model.init(key, dummy_image_batch)
    params = variables['params']

import jax
import jax.numpy as jnp
from model import CNN

def main():
    # Test CNN model initialization
    cnn_model = CNN(num_classes=10)
    key = jax.random.key(123)
    
    # Create dummy batch for initialization
    dummy_batch_size = 1
    dummy_image_batch = jnp.ones((dummy_batch_size, 28, 28, 1), dtype=jnp.float32)
    print(f"Dummy image batch shape: {dummy_image_batch.shape}")

    # Initialize parameters
    variables = cnn_model.init(key, dummy_image_batch)
    params = variables['params']

Python

print("CNN parameters initialized")
print(f"Conv1 kernel shape: {params['conv1']['kernel'].shape}")

# Test forward pass
output_logits = cnn_model.apply({'params': params}, dummy_image_batch)
print(f"Output logits shape: {output_logits.shape}")
assert output_logits.shape == (dummy_batch_size, 10), "Output shape mismatch!"

print("CNN parameters initialized")
print(f"Conv1 kernel shape: {params['conv1']['kernel'].shape}")

# Test forward pass
output_logits = cnn_model.apply({'params': params}, dummy_image_batch)
print(f"Output logits shape: {output_logits.shape}")
assert output_logits.shape == (dummy_batch_size, 10), "Output shape mismatch!"

text

Dummy image batch shape: (1, 28, 28, 1)
CNN parameters initialized
Conv1 kernel shape: (3, 3, 1, 32)
Output logits shape: (1, 10)

Dummy image batch shape: (1, 28, 28, 1)
CNN parameters initialized
Conv1 kernel shape: (3, 3, 1, 32)
Output logits shape: (1, 10)

Python

import flax.linen as nn
import jax.numpy as jnp

class CNN(nn.Module):
    num_classes: int = 10

    @nn.compact
    def __call__(self, x):
        # ...

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

Python

# Flatten for dense layers
x = jnp.reshape(x, (x.shape[0], -1))

# Dense layer
x = nn.Dense(features=128, name='dense1')(x)
x = nn.relu(x)

# Output layer (logits)
x = nn.Dense(features=self.num_classes, name='output_logits')(x)

Python

import jax
import jax.numpy as jnp
from model import CNN

def main():
    # Test CNN model initialization
    cnn_model = CNN(num_classes=10)
    key = jax.random.key(123)
    
    # Create dummy batch for initialization
    dummy_batch_size = 1
    dummy_image_batch = jnp.ones((dummy_batch_size, 28, 28, 1), dtype=jnp.float32)
    print(f"Dummy image batch shape: {dummy_image_batch.shape}")

    # Initialize parameters
    variables = cnn_model.init(key, dummy_image_batch)
    params = variables['params']

Python

print("CNN parameters initialized")
print(f"Conv1 kernel shape: {params['conv1']['kernel'].shape}")

# Test forward pass
output_logits = cnn_model.apply({'params': params}, dummy_image_batch)
print(f"Output logits shape: {output_logits.shape}")
assert output_logits.shape == (dummy_batch_size, 10), "Output shape mismatch!"

text

Dummy image batch shape: (1, 28, 28, 1)
CNN parameters initialized
Conv1 kernel shape: (3, 3, 1, 32)
Output logits shape: (1, 10)

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