Skip to content

← Flax NNX step 47 of 100

Hard Primitives

NNX ResNet Basic Block

Jump to the task

Why this matters

A residual block adds a learned transformation to its input. The transformation can adjust the existing features instead of replacing them completely. This exercise builds a two-convolution block with an identity shortcut and a ReLU after the addition.

If the residual branch returns zero, this block returns relu(x), not necessarily x. The final ReLU still removes negative values. The shortcut provides a direct term in the derivative of the sum, but the final ReLU can block gradients where that sum is negative. Residual connections do not guarantee that every gradient survives.

Follow the two paths

The residual branch computes:

x -> conv1 -> bn1 -> ReLU -> conv2 -> bn2

The shortcut carries x unchanged. Add the two paths, then apply ReLU: out = jax.nn.relu(residual + x). Do not apply another ReLU between bn2 and the addition. Negative residual values must be able to reduce the input before the final activation.

Both convolutions use a 3-by-3 kernel, stride 1, and SAME padding. They map C channels to C channels, so all intermediate arrays and the output have shape (height, width, C). SAME padding supplies zeros around the spatial boundaries here.

An identity shortcut requires the full input and output shapes to match. This exercise keeps both spatial dimensions and channel count fixed. Other residual blocks can change those dimensions and use a projection shortcut; the name "basic block" does not rule that out.

What BatchNorm does here

Each nnx.BatchNorm treats the last axis as channels. The input has no batch axis, so training-mode statistics are computed over the height and width axes of this one image, independently per channel. It is not LayerNorm over the channels of each pixel.

With use_running_average=False, each BatchNorm uses the current activation's channel mean and variance for its output. It also updates its stored running statistics. These are separate operations: the output does not use the newly updated running averages.

Use the default settings explicitly: epsilon=1e-5, momentum=0.99, trainable scale initialized to one and bias initialized to zero. Running means start at zero and running variances at one. For current channel statistics mean and var, the stored updates are:

running_mean = 0.99 * running_mean + 0.01 * mean
running_var  = 0.99 * running_var  + 0.01 * var

The variance uses the population convention, without an unbiased sample correction. NNX's default fast variance calculation clamps mean(h * h) - mean(h) ** 2 at zero for numerical roundoff.

With use_running_average=True, normalization instead reads the stored statistics and leaves them unchanged. NNX keeps this state in the BatchNorm submodules; no separate mutable collection argument is needed at the call site.

A small state example

This complete example has one channel and two spatial positions. Its training mean is 2 and variance is 1:

import jax.numpy as jnp
import numpy as np
from flax import nnx

bn = nnx.BatchNorm(1, epsilon=1e-5, momentum=0.99, rngs=nnx.Rngs(0))
x = jnp.array([[[1.], [3.]]])
y = bn(x, use_running_average=False)
np.testing.assert_allclose(
    y.reshape(-1), jnp.array([-1., 1.]) / jnp.sqrt(1. + 1e-5),
    rtol=1e-5,
)
np.testing.assert_allclose(bn.mean[...], [0.02], rtol=1e-5)
np.testing.assert_allclose(bn.var[...], [1.], rtol=1e-5)

With only one spatial position, the training variance is zero and the centered activation is zero. The epsilon keeps normalization defined; with the initial BatchNorm biases, this block then returns relu(x).

Problem

Write resnet_basic_forward(seed, x, features) using a BasicBlock(nnx.Module).

  • Input x has shape (height, width, C) with positive dimensions. C equals int(features).
  • Create one nnx.Rngs(int(seed)). Construct conv1, bn1, conv2, then bn2 in that order, passing the same RNG container to each.
  • Each convolution is nnx.Conv(C, C, kernel_size=(3, 3), strides=1, padding="SAME", rngs=rngs). Keep its default initializer and enabled bias, even though BatchNorm follows it.
  • Each BatchNorm has C features and the settings described above.
  • Apply conv1, bn1, ReLU, conv2, bn2, add the original x, then ReLU. Pass use_running_average=False to both BatchNorm calls.
  • Return out.reshape(-1), of length height * width * C.

The entrypoint creates a fresh block for every call. This is one training-mode forward pass, not a parameter-training step: it updates running statistics but performs no gradient calculation or optimizer update. Those statistics are not returned or retained between calls to the entrypoint.