Skip to content

← Flax step 51 of 100

Hard Primitives

Tiny ResNet Classifier

Jump to the task

Turn spatial features into class scores

This small classifier has three parts. A convolutional stem creates features, two residual blocks refine them, and a pooled linear head produces one score per class. The scores are logits, not probabilities. The exercise initializes and evaluates the network; it does not train a classifier or measure classification accuracy.

All convolutional layers preserve height and width. This is a compact teaching network, without the downsampling stages and changing widths found in larger ResNet architectures.

Track the shapes

Let the image have shape (H, W, C) and let K = num_classes.

Stage Operation Output shape
Input Add a batch axis (1, H, W, C)
Stem Conv3x3(C), BatchNorm, ReLU (1, H, W, C)
Block 1 Shape-preserving basic block (1, H, W, C)
Block 2 A second basic block with its own parameters (1, H, W, C)
Pooling Mean over axes (1, 2) (1, C)
Head Dense(K) (1, K)
Return Flatten logits (K,)

This task sets the stem width to the image's channel count. That is a specific architecture choice, not a requirement for identity shortcuts. A stem could map an RGB image to a different width as long as the subsequent blocks preserve that new width.

Each basic block computes relu(F(h) + h), with F = Conv3x3 -> BatchNorm -> ReLU -> Conv3x3 -> BatchNorm. Both blocks use identity shortcuts. Their weights are independent: create two module instances rather than applying one shared block twice.

What global average pooling preserves

Pooling keeps one mean per channel and removes explicit spatial coordinates. For a batched feature map, average axes (1, 2) only. Keep the batch and channel axes. The head then needs C*K + K parameters, compared with H*W*C*K + K for flattening all spatial positions before Dense.

The head's parameter shape does not depend on image resolution, but its predictions and normalization statistics can still change with resolution. Averaging is invariant to rearranging positions in an already-computed feature map; that does not make the whole convolutional network invariant to arbitrary rearrangements of image pixels.

A Dense layer on the last axis of a feature map is a shared per-position affine map, not a flatten operation. With the same weights, applying it before spatial averaging gives the same mathematical result as applying it afterward. Pooling first avoids evaluating that affine map at every position. Here is a small check:

import jax.numpy as jnp
import numpy as np

h = jnp.array([1., 2., 3., 4., 5., 6., 7., 8.]).reshape(1, 2, 2, 2)
w = jnp.array([[1., -1.], [2., 0.5]])
b = jnp.array([0.25, -0.5])
after_pool = h.mean(axis=(1, 2)) @ w + b
before_pool = (h @ w + b).mean(axis=(1, 2))
np.testing.assert_allclose(after_pool, [[14.25, -2.]])
np.testing.assert_allclose(before_pool, after_pool)

This equivalence does not generally hold if you insert a nonlinear activation between the per-position Dense and pooling.

Problem

Implement resnet_classifier_forward(seed, image, num_classes):

  1. Define a compact BasicBlock(features) with the exact two-convolution order described above. Use stride-one 3-by-3 convolutions with SAME padding and nn.BatchNorm(use_running_average=False).
  2. Define a compact TinyResNet(stem_features, num_classes). Its stem is a 3-by-3 SAME convolution at stride one, training-mode BatchNorm, and ReLU. Follow it with two separate BasicBlock(stem_features) instances, mean pooling over axes (1, 2), and nn.Dense(num_classes).
  3. Set stem_features = image.shape[-1]. Initialize TinyResNet directly with jax.random.PRNGKey(int(seed)) and image[None, ...].
  4. Apply it to the same batched input with the complete variables tree and mutable=["batch_stats"]. Unpack the output and updated state.
  5. Return the logits flattened in row-major order, with no softmax.

Keep all default initializers, convolution biases, Dense bias, and BatchNorm settings. Leave child modules unnamed. At the classifier level this gives Conv_0, BatchNorm_0, BasicBlock_0, BasicBlock_1, and Dense_0. Each block has its own nested Conv and BatchNorm names. Do not split the initialization key or change this module nesting.

BatchNorm computes current-input statistics over batch and spatial positions for each channel. The call returns updated running statistics separately; the wrapper discards them. It is not an evaluation-mode pass using learned running averages, despite producing classifier logits.

Inputs: image is a float array (H, W, C) with positive dimensions. seed and the positive num_classes are integer-valued scalars, possibly floats; convert them to Python integers. Output: a one-dimensional array of shape (num_classes,). Keep both scalar setup arguments static when JIT-compiling the wrapper.