Skip to content

← Flax step 52 of 100

Hard Primitives

Tiny U-Net

Jump to the task

Combine coarse and fine feature maps

This tiny U-Net has a contracting encoder and an expanding decoder. The encoder pools spatial positions to build lower-resolution features. The decoder upsamples, then joins those features with saved encoder maps at the matching resolution. These skip connections give the decoder access to features that did not pass through the smallest spatial map.

The skips help preserve access to spatial detail, but do not guarantee reconstruction or accurate segmentation. This exercise performs one randomly initialized forward pass. Its output is a feature map with base_features channels, not class probabilities, segmentation labels, or an image reconstructed by an inverse operation.

Read the architecture by shape

Let f = base_features. Add a batch axis to the image before entering the model. For a (4, 4, 1) image and f=2:

Name Operation Shape
image Add batch axis (1, 4, 4, 1)
e1 ConvBlock(f) (1, 4, 4, 2)
p1 MaxPool2x2, stride 2 (1, 2, 2, 2)
e2 ConvBlock(2f) (1, 2, 2, 4)
p2 MaxPool2x2, stride 2 (1, 1, 1, 4)
b ConvBlock(4f) (1, 1, 1, 8)
u2 ConvTranspose2x2(2f), stride 2 (1, 2, 2, 4)
joined2 Concatenate [u2, e2] (1, 2, 2, 8)
d2 ConvBlock(2f) (1, 2, 2, 4)
u1 ConvTranspose2x2(f), stride 2 (1, 4, 4, 2)
joined1 Concatenate [u1, e1] (1, 4, 4, 4)
d1 ConvBlock(f) (1, 4, 4, 2)

Each ConvBlock is one stride-one 3-by-3 SAME convolution, followed by training-mode BatchNorm and ReLU. It is not a two-convolution block. There is no final class projection in this task.

Concatenate in the specified order

Concatenation along axis=-1 keeps all channels from both inputs. Its channel count is the sum of their widths. The widths happen to match at both joins in this model, so they double. This differs from an elementwise residual addition, which keeps the original width.

The order matters to the next convolution: use decoder channels first, then encoder channels. Here is a small example of that ordering:

import jax.numpy as jnp
import numpy as np

upsampled = jnp.array([1., 2.]).reshape(1, 1, 1, 2)
encoder_skip = jnp.array([10., 20.]).reshape(1, 1, 1, 2)
joined = jnp.concatenate([upsampled, encoder_skip], axis=-1)
np.testing.assert_array_equal(joined.reshape(-1), [1., 2., 10., 20.])
assert joined.shape == (1, 1, 1, 4)

Match spatial dimensions

Both pooling layers use a 2-by-2 window, stride two, and VALID padding. Both transposed convolutions use a 2-by-2 kernel, stride two, and VALID padding. With these settings, each pooling stage halves an even input dimension and each upsampling stage doubles it.

Height and width must be positive multiples of four. They need not be equal or powers of two: a 4-by-12 image satisfies the requirement. Arbitrary odd dimensions would need a deliberate padding or cropping policy before the skip joins; do not invent one for this task. Transposed convolution is a learned upsampling operation, not an inverse that recovers values removed by max pooling.

Problem

Implement unet_forward(seed, image, base_features):

  1. Define a compact ConvBlock(features) with nn.Conv(features, (3, 3), padding="SAME"), nn.BatchNorm(use_running_average=False), then nn.relu.
  2. Define a compact TinyUNet(base_features). Build the encoder, bottom, and decoder in the table's order. Save e1 and e2 before their corresponding pooling operations.
  3. Use nn.max_pool with window_shape=(2, 2), strides=(2, 2), and padding="VALID". Use nn.ConvTranspose with kernel_size=(2, 2), strides=(2, 2), padding="VALID", and the default transpose_kernel=False.
  4. Concatenate [u2, e2] and [u1, e1] on the channel axis before the decoder ConvBlocks. Return the final ConvBlock's feature map.
  5. Initialize TinyUNet directly on image[None, ...] with jax.random.PRNGKey(int(seed)). Apply it to that input with the complete variables tree and mutable=["batch_stats"].
  6. Unpack the output and updated state, and return out.reshape(-1).

Keep all default biases, initializers, and BatchNorm settings. Leave child modules unnamed. The five ConvBlocks are ConvBlock_0 through ConvBlock_4; the upsamplers are ConvTranspose_0 and ConvTranspose_1. Each ConvBlock has its own Conv_0 and BatchNorm_0. Do not change the nesting or split the initialization key.

Inputs: image is a float array (H, W, C) with positive C and positive spatial dimensions divisible by four. base_features is a positive integer-valued scalar and seed is an integer-valued scalar; convert scalar inputs to Python integers. Output: a one-dimensional array of length H * W * base_features.

BatchNorm uses current batch and spatial statistics, returning updated running statistics separately. The wrapper discards these updates. With a 4-by-4 image, the bottom has one spatial position and one batch element, so training-mode BatchNorm centers its only value per channel to zero. At initialization its offset is also zero. The skip paths still carry features from larger maps; this small test case does not demonstrate a rich learned bottleneck. Keep both scalar setup arguments static when JIT-compiling the wrapper.