Medium End-to-End

Train Image Classifier End-to-End

Implement train_image_classifier: train a small convolutional classifier using manually computed gradients. Each epoch makes one update from the entire supplied batch. Carry both sets of updated weights into the next epoch.

The model applies a convolution, ReLU, max pooling and a linear classifier. Train both the convolution filters and the classifier matrix, but return only the final classifier matrix. Filter updates affect the features used by later epochs, so they matter even though filters are not returned.

Inputs and return value

x: (N, 1, H, W), a batch of grayscale images.
y: (N,), integer-valued class labels in [0, K), delivered as floats.
w_conv: (C_out, 1, kH, kW), the convolution filters.
w_fc: (C_outH_poolW_pool, K), the classifier matrix.
lr: a finite, nonnegative learning rate.
n_epochs: a nonnegative integer, the number of full-batch updates.

Here H_conv = H-kH+1, W_conv = W-kW+1, H_pool = H_conv//2 and W_pool = W_conv//2. Dimensions are positive, H_conv and W_conv are at least two, and all shapes are consistent. The floating-point tensors are finite and share a dtype and device. Cast labels to integers for indexing.

Return the final w_fc with its original shape, dtype and device. Leave all inputs unchanged. If n_epochs=0 or lr=0, return the original w_fc values.

Forward pass for each epoch

Use the current weights to compute:

conv = conv2d(x, w_conv)
relu = maximum(conv, 0)
pool = max_pool_2x2_stride2(relu)
flat = pool.reshape(N, -1)
logits = flat @ w_fc
probs = softmax(logits, axis=-1)
loss = -sum(log_softmax(logits)[n, y[n]] for n in range(N)) / N

The convolution has stride one, no padding and no bias. Use each filter as stored, without flipping it: conv[n,c,i,j] is the sum of x[n,0,i+kh,j+kw] * w_conv[c,0,kh,kw] over kh and kw.

Pool each channel separately with non-overlapping 2-by-2 windows. Start at the top-left corner. Discard a trailing row or column if the convolution's spatial dimension is odd. When a window has tied maxima, choose its first maximum in row-major order: top-left, top-right, bottom-left, bottom-right. Save that position for the backward pass.

Flatten pooled features in channel, row, then column order. The classifier has no bias. For stable softmax, subtract each row's maximum before exponentiating and normalizing. The loss averages over examples only.

Manual backward pass

Start from the mean-loss gradient and continue through the classifier:

dlogits = (probs - one_hot(y, K)) / N
dw_fc = flat.T @ dlogits
dflat = dlogits @ w_fc.T
dpool = dflat.reshape(N, C_out, H_pool, W_pool)

For each pooling window, send its incoming gradient only to the saved maximum position. All other positions receive zero, including any discarded trailing row or column. Then pass gradients through ReLU only where conv > 0; use a zero derivative at conv=0.

Accumulate filter gradients over examples and convolution output positions:

dw_conv[c,0,kh,kw] = sum(
    x[n,0,i+kh,j+kw] * dconv[n,c,i,j]
    for n, i, j over their valid ranges
)

Do not divide this sum by N again. The factor 1/N entered at dlogits and has already propagated through the rest of the backward pass.

Compute both weight gradients using the weights from the start of the epoch. Then update both:

w_conv_new = w_conv - lr * dw_conv
w_fc_new = w_fc - lr * dw_fc

Use these new weights for the next epoch's forward pass. Do not shuffle the batch or introduce mini-batches, momentum or weight decay.

Build the layers and derivatives from basic tensor operations in either framework. Do not use packaged convolution, pooling or neural-network layers, automatic differentiation, backward(), jax.grad, or optimizer helpers. Store the forward values needed by the backward pass.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Apply the same filter at many locations

A convolutional layer slides a learned kernel over its input. At each location, it multiplies the covered values by the kernel weights and sums them. Deep-learning libraries usually implement cross-correlation under the name convolution: the kernel is not reversed.

For input [1, 2, 3] and kernel [1, -1], with no padding and stride one, the outputs are 1*1 + 2*(-1) = -1 and 2*1 + 3*(-1) = -1.

Work out the output shape

For one spatial axis, the output length is

$$\left\lfloor\frac{L+2P-D(K-1)-1}{S}\right\rfloor+1.$$

L is the input length, K the kernel size, P the padding on each end, S the stride between windows, and D the spacing between kernel elements (dilation). For a 2D input, apply the formula separately to height and width. Check the channel layout too: an ordinary convolution sums over both spatial positions in the window and input channels.

A depthwise convolution filters each input channel separately. A following pointwise convolution, using a 1-by-1 spatial kernel, mixes channels. This uses a different set of parameters from a full spatial convolution; the two constructions are not generally interchangeable.

Pool nearby features

Max pooling keeps the largest value in each window. Average pooling takes the window's mean. During backpropagation, max pooling sends the gradient to the selected position, while average pooling distributes it across the values that contributed to the mean. Overlapping windows can contribute more than once to the same input position, so add those gradients.

Follow the exercise's rules for tied maxima and whether padding counts in an average. For a small CNN, also save the forward-pass values needed for activation and pooling derivatives. Add each shared kernel's gradient contributions across all locations, then update weights after computing all gradients.

Detecting runtime…