Skip to content

← Flax step 49 of 100

Hard Primitives

ResNet Basic Block

Jump to the task

Learn a correction to the input

This block computes relu(F(x) + x). The residual branch F contains two convolutions with BatchNorm, while the shortcut carries the input directly to the addition. Both branches have the same shape.

The shortcut provides an additional path for gradients, but does not guarantee that gradients stay unchanged or that training succeeds. The final ReLU can block gradients where the summed value is negative. Also, setting F(x) to zero gives relu(x), which equals x only where x is nonnegative. For example, [-2, 1] becomes [0, 1].

Follow the operation order

residual branch: x -> Conv3x3 -> BatchNorm -> ReLU -> Conv3x3 -> BatchNorm
shortcut:        x --------------------------------------------------+
                                                                      |
                                                 add both -> ReLU -> output

There is no activation between the second BatchNorm and the addition. Adding one would change the residual function; it is not what defines a pre-activation ResNet block. This exercise uses the order above and no projection on the shortcut.

Both convolutions use stride one and SAME padding. The input channel count must equal features, so height, width, and channels all match at the addition. Blocks that change these dimensions need a compatible shortcut, but that is outside this task.

BatchNorm uses spatial positions too

Add a batch axis to the (H, W, C) input, giving (1, H, W, C). With the default last feature axis, BatchNorm computes a mean and variance for each channel over axes (0, 1, 2): batch, height, and width. A batch size of one still supplies H * W values per channel. Do not normalize each pixel across channels as LayerNorm would.

Use nn.BatchNorm(use_running_average=False) in both layers. The output uses current-input statistics, not the stored running averages. Each BatchNorm has trainable scale and bias in params, plus running mean and variance in batch_stats.

Flax returns updated state separately from apply; it does not modify the caller's variables dictionary in place. This small example shows the distinction:

import jax
import jax.numpy as jnp
import flax.linen as nn
import numpy as np

x = jnp.array([1., 5.]).reshape(1, 1, 2, 1)
bn = nn.BatchNorm(use_running_average=False)
variables = bn.init(jax.random.PRNGKey(0), x)
y, updates = bn.apply(variables, x, mutable=["batch_stats"])
np.testing.assert_allclose(y.reshape(-1), [-2 / np.sqrt(4.00001), 2 / np.sqrt(4.00001)])
np.testing.assert_allclose(variables["batch_stats"]["mean"], [0.])
np.testing.assert_allclose(updates["batch_stats"]["mean"], [0.03], atol=1e-6)

This uses the default epsilon 1e-5 and momentum 0.99. In a training loop you would carry returned statistics into later calls. Here the wrapper returns only the block output and discards those updates. Applying this training-mode BatchNorm without allowing batch_stats updates raises an error.

Problem

Implement resnet_basic_forward(seed, x, features):

  1. Define a compact BasicBlock(nn.Module) with an integer features field. Save its input as residual.
  2. Apply nn.Conv(features, (3, 3), strides=(1, 1), padding="SAME"), then nn.BatchNorm(use_running_average=False), then nn.relu.
  3. Apply a second convolution and BatchNorm with the same settings. Add residual, then apply the final nn.relu.
  4. Instantiate this block directly, initialize it with jax.random.PRNGKey(int(seed)) and x[None, ...], and apply it to that same batched input with mutable=["batch_stats"].
  5. Unpack (out, updated_state) and return out.reshape(-1).

Keep the default Conv biases and initializers and the default BatchNorm settings. Both Conv biases are included in this particular exercise. Leave the child modules unnamed so they become Conv_0, BatchNorm_0, Conv_1, and BatchNorm_1 in creation order. Changing parameter paths or splitting the initialization key changes the seeded result.

Inputs: x is a float array (H, W, C) with positive dimensions and C == int(features). seed and features are integer-valued scalars, possibly supplied as floats; convert them to Python integers. Output: a one-dimensional array of length H * W * C, in row-major order. This is one freshly initialized forward pass, not a training step.

Pass both params and batch_stats to apply, or pass the complete variables tree returned by init. Keep seed and features static when JIT-compiling the wrapper. With one spatial position, each channel has only one value for normalization; do not substitute running-statistic mode to avoid that boundary case.