NNX Tiny ResNet Classifier

Why this matters

A residual block transforms a feature map. An image classifier also needs a way to create those features from pixels and turn them into one prediction per image. This exercise combines an initial convolution, two residual blocks, spatial averaging, and a linear classification head.

Build a fresh model and return its logits. There is no training or checkpoint loading, so these are not yet useful predictions. The four-channel body and two blocks define this small exercise; it is not a standard ResNet-18 or ResNet-50 architecture.

From pixels to features

Input image has shape (height, width, C), without a batch axis. The initial layer, often called the stem, maps C image channels to four feature channels:

image -> 3-by-3 convolution -> BatchNorm -> ReLU

Use stride 1 and SAME padding. The result has shape (height, width, 4). Padding adds zeros around the spatial boundaries; this model does not downsample. The stem has no shortcut, even when the image happens to have four channels. That is an architectural choice, not a consequence of mismatched channel counts.

Two residual blocks

Create two distinct blocks, each with its own parameters and BatchNorm state. Each block takes a four-channel feature map x and computes a residual branch:

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

Both convolutions map four channels to four channels with a 3-by-3 kernel, stride 1, and SAME padding. Return jax.nn.relu(residual + x). Keep the second ReLU after the addition; applying it before the addition would prevent negative residual values from reducing x. The shortcut is the unchanged input because the full shapes match.

Apply block1, then block2. Calling one block twice would share its parameters and running statistics, which is a different model.

BatchNorm during this forward pass

Use use_running_average=False in the stem and both blocks, for five BatchNorm calls in total. Each channel's mean and population variance are computed over height and width of this one image. BatchNorm normalizes with those current statistics and also updates its stored running statistics.

Use epsilon=1e-5 and momentum=0.99, with default scale one, bias zero, running mean zero and running variance one. Each stored mean or variance becomes 0.99 * old + 0.01 * current. Evaluation with use_running_average=True instead reads stored statistics without changing them.

A training-mode forward pass is not an optimizer step. It updates running statistics but does not change convolution kernels, linear weights, or BatchNorm scale and bias. The entrypoint creates a new model every call, so it does not retain these statistics across calls.

Average spatial positions, then classify

After the blocks, compute jnp.mean(x, axis=(0, 1)). This reduces (height, width, 4) to a four-element vector, retaining one average per channel. Apply nnx.Linear(4, K) to obtain K logits.

This complete example isolates the pooling and head with known values:

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

features = jnp.array([[[1., 2., 3., 4.], [3., 4., 5., 6.]]])
pooled = features.mean(axis=(0, 1))
head = nnx.Linear(4, 2, rngs=nnx.Rngs(0))
head.kernel[...] = jnp.array([[1., 0.], [0., 1.],
                             [1., 0.], [0., 1.]])
head.bias[...] = jnp.array([0.5, -0.5])
np.testing.assert_allclose(pooled, [2., 3., 4., 5.])
np.testing.assert_allclose(head(pooled), [6.5, 7.5])

The head has 4 * K + K parameters, including its bias, regardless of image dimensions. Flattening the entire feature map before a dense head would require height * width * 4 * K + K parameters and fix its expected input size.

Averaging is invariant to reordering the final feature map's spatial positions. That does not guarantee that shifting the input image leaves the whole model unchanged: convolution boundaries and padding can change the computed features. Likewise, accepting different image sizes does not mean predictions are independent of resolution.