Skip to content

← Flax step 50 of 100

Hard Primitives

ResNet Bottleneck Block

Jump to the task

Do spatial filtering at an intermediate width

This residual block uses three convolutions: a 1-by-1 channel projection, a 3-by-3 spatial filter, and a final 1-by-1 projection. If the intermediate width is smaller than the input width, the spatial filter operates on fewer channels. The final projection restores the channel count needed for the shortcut addition.

Let O = out_features and M = mid_features. The task requires x.shape[-1] == O, so the channel sequence is O -> M -> M -> O. The input is also carried unchanged along an identity shortcut. A smaller M can reduce the weight count, but does not guarantee the same representational capacity, accuracy, or runtime as a wider block. Use the supplied M; do not assume it is always O // 4.

Follow the three stages

Stage Operation Output shape, including batch
Input Add one batch axis (1, H, W, O)
First projection Conv1x1, BatchNorm, ReLU (1, H, W, M)
Spatial filtering Conv3x3, BatchNorm, ReLU (1, H, W, M)
Final projection Conv1x1, BatchNorm (1, H, W, O)
Residual output Add original input, then ReLU (1, H, W, O)

All convolutions use stride one and SAME padding. A 1-by-1 convolution mixes channels at each position without reading neighboring positions. The 3-by-3 convolution also mixes nearby positions.

Do not insert a ReLU after the third BatchNorm before the addition. That would change the requested function. There is no shortcut projection or spatial downsampling in this exercise.

Compare complete weight counts

The three kernels contain O*M + 9*M*M + M*O scalar weights. For O=256 and M=64, this is 69,632. Two full-width 3-by-3 kernels would contain 2 * 9 * 256 * 256 = 1,179,648 weights, about 16.9 times as many. These counts exclude biases and BatchNorm.

This implementation keeps all convolution biases and both trainable BatchNorm parameters. Those add 3 * (2*M + O) scalars: one bias, one scale, and one normalization offset for every intermediate/output channel. The running means and variances are state, not trainable parameters, and add another 2 * (2*M + O) stored values.

For the smaller example O=4, M=2, there are 52 kernel weights, 24 other trainable scalars, and 16 running-statistic values. This makes a useful check of the initialized variables tree.

BatchNorm and the shortcut

Each BatchNorm uses use_running_average=False. With channels-last batched input, it normalizes each channel over batch, height, and width. Keep its default epsilon, momentum, scale, and bias settings. Training-mode application returns updated running statistics separately when mutable=["batch_stats"] is set; it does not change the supplied dictionary in place. The wrapper discards those returned statistics because the requested result is only this forward pass.

A residual connection does not guarantee unchanged gradients. The output is relu(F(x) + x), so even a zero residual branch produces relu(x), not an identity map for negative inputs. The final ReLU can also block gradients where the sum is negative.

Problem

Implement resnet_bottleneck_forward(seed, x, mid_features, out_features):

  1. Define a compact BottleneckBlock(nn.Module) with integer fields mid_features and out_features. Save its input as residual.
  2. Apply nn.Conv(mid_features, (1, 1), padding="SAME"), nn.BatchNorm(use_running_average=False), then nn.relu.
  3. Apply nn.Conv(mid_features, (3, 3), padding="SAME"), another training-mode BatchNorm, then nn.relu.
  4. Apply nn.Conv(out_features, (1, 1), padding="SAME"), then a third training-mode BatchNorm. Add residual and apply ReLU.
  5. Initialize the block directly with jax.random.PRNGKey(int(seed)) and x[None, ...]. Apply it to that input with the complete variables tree and mutable=["batch_stats"].
  6. Unpack the output and updated state, then return out.reshape(-1).

Keep all Conv and BatchNorm defaults, including convolution biases. Leave child modules unnamed: creation order produces Conv_0 through Conv_2 and BatchNorm_0 through BatchNorm_2. Do not split the initialization key or add a parent module, which would change the seeded parameter paths.

Inputs: x is a float array (H, W, O) with positive dimensions. mid_features and out_features are positive integer-valued scalars, and seed is an integer-valued scalar. Scalars may arrive as floats; convert them to Python integers. Output: a one-dimensional array of length H * W * O in row-major order. This is a freshly initialized forward pass, with no optimizer step.

Keep all scalar setup arguments static if JIT-compiling the wrapper. Check the entire shape before residual addition: broadcasting can hide some shape mistakes rather than raising an error.