Skip to content

← Flax NNX step 49 of 100

Hard Primitives

NNX Tiny U-Net

Jump to the task

Why this matters

An encoder reduces spatial resolution to process features at a coarser scale. A decoder increases resolution again. U-Net connects these paths by passing encoder features directly to decoder stages at matching resolutions. Those skip connections give the decoder access to detail that downsampling may discard.

This exercise builds a small two-level encoder-decoder with that pattern. It returns a feature map, not segmentation labels or class probabilities. There is no final prediction head, training step, or pretrained checkpoint. Each call runs a freshly initialized model.

Shapes through the network

Input image has shape (H, W, C), without a batch axis. Require positive channel counts and H and W to be positive multiples of four. This ensures the two downsampling and upsampling stages line up exactly, without cropping or padding the skip tensors.

Let f = int(base_features), a positive integer. A Conv-BN-ReLU block applies a 3-by-3 convolution with stride 1 and SAME padding, then BatchNorm and ReLU. It preserves spatial dimensions.

Stage Operation Output shape
enc1 Conv-BN-ReLU, C to f channels (H, W, f)
pool1 2-by-2 max pool, stride 2 (H/2, W/2, f)
enc2 Conv-BN-ReLU, f to 2f (H/2, W/2, 2f)
pool2 2-by-2 max pool, stride 2 (H/4, W/4, 2f)
bottom Conv-BN-ReLU, 2f to 4f (H/4, W/4, 4f)
up2 Transposed convolution, 4f to 2f (H/2, W/2, 2f)
concat2 Concatenate up2 output, then enc2 output (H/2, W/2, 4f)
dec2 Conv-BN-ReLU, 4f to 2f (H/2, W/2, 2f)
up1 Transposed convolution, 2f to f (H, W, f)
concat1 Concatenate up1 output, then enc1 output (H, W, 2f)
dec1 Conv-BN-ReLU, 2f to f (H, W, f)

Save the outputs of enc1 and enc2 before pooling. These are the skip tensors. For nonmultiples of four, VALID pooling drops trailing rows or columns; simply doubling the sizes later will not recover all the dimensions needed for these concatenations.

Pool and upsample

Implement pooling with lax.reduce_window(x, -jnp.inf, lax.max, (2, 2, 1), (2, 2, 1), "VALID"). The final window size and stride are one, so channels stay separate. Negative infinity is the identity for a maximum: it cannot replace a finite value in a nonempty window. Assume finite image inputs.

Each upsampling layer is nnx.ConvTranspose with kernel_size=(2, 2), strides=2, and padding="VALID". With these settings it doubles both spatial dimensions. A transposed convolution is a learned operation, not an inverse that recovers pixels discarded by pooling. Keep the default transpose_kernel=False convention and default initialization used by NNX.

Concatenate the skip features

Use jnp.concatenate([upsampled, skip], axis=-1). This joins channels while preserving height and width. The order matters because the next convolution assigns different weights to different input channels. The decoder block therefore needs twice the upsampling layer's output channel count.

This complete example isolates pooling and channel concatenation. The zero tensor stands in for an upsampling result; it is not a transposed-convolution implementation:

import jax.numpy as jnp
import numpy as np
from jax import lax

skip = jnp.arange(1., 17.).reshape(4, 4, 1)
pooled = lax.reduce_window(
    skip, -jnp.inf, lax.max, (2, 2, 1), (2, 2, 1), "VALID"
)
np.testing.assert_allclose(pooled[..., 0], [[6., 8.], [14., 16.]])
upsampled = jnp.zeros_like(skip)
joined = jnp.concatenate([upsampled, skip], axis=-1)
assert joined.shape == (4, 4, 2)
np.testing.assert_array_equal(joined[..., :1], upsampled)
np.testing.assert_array_equal(joined[..., 1:], skip)

Addition would combine corresponding channels and retain width f. Concatenation retains the two sets of features separately for the next convolution to mix. This task requires concatenation; neither operation is universally better for every architecture.

BatchNorm and the small-image case

Pass use_running_average=False to all five Conv-BN-ReLU blocks. Each BatchNorm computes per-channel statistics over the height and width of its current feature map. Use epsilon=1e-5, momentum=0.99, and default scale, bias and running-statistic initialization.

The forward pass uses current statistics and updates stored means and variances, but it does not optimize weights. The entrypoint creates a new model each call, so those statistics are not retained between calls. Evaluation mode would instead read stored statistics.

For a 4-by-4 image, the bottom map is 1-by-1. Its training-mode BatchNorm centers each single channel value to zero; with the initial bias, the bottom output is zero. The saved encoder features still provide input to the decoder through the skip connections. The bottleneck does not carry useful image information in this special initialized case.

Problem

Write unet_forward(seed, image, base_features) with NNX modules. Cast seed and base_features to integers. Use one nnx.Rngs(int(seed)) and construct modules in this order:

  1. enc1, enc2, bottom.
  2. up2, dec2, up1, dec1.

Within each Conv-BN-ReLU block, construct the convolution before BatchNorm. Use the channel widths and spatial settings in the table. Keep default NNX kernel initializers and enabled biases for both ordinary and transposed convolutions.

Follow the table's forward order, retaining both skip tensors and concatenating upsampled features first. Return out.reshape(-1) from dec1, with shape (H * W * base_features,). The final ReLU makes these feature values nonnegative. Do not add a classification head, sigmoid, softmax, or spatial averaging.