Simple CNN — MNIST

Every architecture diagram in this gallery is really just a sequence of convolutions with a shape attached to each step. This model is small enough that you can hold the entire sequence in your head at once — six blocks, every one of them a 3×3 or 1×1 convolution, each shrinking or reshaping the volume by a fixed, computable amount. That makes it the right place to build the habit this page is for: at every step, ask what is the shape now, and how many numbers does this step cost.

The input is a single 28×28 grayscale image — one channel, so the input volume is $28 \times 28 \times 1$. Trace it block by block, and every number below comes straight from the trained checkpoint's layer definitions:

Two shape rules explain every entry in that output shape column, and both are the convolution arithmetic you've seen before, applied with $p=0$. A plain 3×3 convolution with no padding and stride 1 takes $n \times n$ to $(n-2) \times (n-2)$ — each 3×3 block in this network costs exactly 2 pixels off each spatial dimension. A 2×2 max-pool with stride 2 halves the spatial size outright. Run those two rules down the table: $28 \to 26 \to 24$ (two plain convs), $\to 12$ (the pool in transition1, the 1×1 conv that follows it doesn't change spatial size), $\to 10 \to 8$ (two more plain convs). Every number in the output shape column is one of those two operations, nothing else.

What a feature map is, here. After conv_block1, the volume is $26 \times 26 \times 16$ — sixteen separate $26 \times 26$ grids, one per output channel, each the record of where one learned 3×3 pattern fired across the digit. A pixel at position $(i, j)$ in channel 7 is the dot product of that channel's kernel with the $3 \times 3$ patch of the input centred near $(i,j)$ — exactly the arithmetic from the convolution page, just with $C_{\text{in}}=1$ so the kernel volume is $3 \times 3 \times 1$, the simplest possible case. By conv_block4, each of the 32 channels in the $8 \times 8 \times 32$ volume is the result of four convolutions stacked — so a single cell there is a function of a much larger patch of the original 28×28 image than the 3×3 a first-layer cell sees. That growth is the receptive field at work, and it is exactly why the last block is where Grad-CAM looks (see the Interpretability tab).

Now size one layer exactly, the way you would read any PyTorch Conv2d definition. conv_block4 maps $32$ input channels to $32$ output channels with a $3 \times 3$ kernel and bias=False. The parameter-count formula for a convolutional layer is $C_{\text{out}}(k \cdot k \cdot C_{\text{in}} + 1)$ when there's a bias term per output channel — drop the $+1$ here because this network's convolutions carry no bias (BatchNorm supplies the shift instead):

That single block is the single most expensive layer in the network: 9,216 of the model's roughly 60,000 total parameters, just under one-sixth, in one 3×3-by-32-by-32 layer. Compare it to transition1's 1×1 convolution — $16 \times (1 \times 1 \times 32) = 512$ — which does real work (it halves the channel count from 32 to 16) for an eighteenth of the cost. That gap is the entire reason 1×1 convolutions exist as a design tool: spatial kernels ($k>1$) are where the parameter budget goes, and $1 \times 1$ convolutions are how architects buy channel-mixing almost for free.

The head: from a volume to ten numbers. After conv_block4, the network holds an $8 \times 8 \times 32$ volume — 2,048 numbers, far too many to feed straight into a 10-way decision without either an enormous dense layer or throwing away spatial information on purpose. This network does the second: global average pooling (GAP) replaces each of the 32 channels' $8 \times 8$ grid with a single number, its mean. The intuition is direct — if a channel is a detector for some learned pattern, its average response across the whole image is a reasonable summary of "how much of this pattern is present, anywhere." GAP turns $8 \times 8 \times 32 \to 1 \times 1 \times 32$, with zero learned parameters — it's an architectural choice, not a layer you train.

The final step is a $1 \times 1$ convolution mapping $32 \to 10$ channels — $10 \times (1 \times 1 \times 32) = 320$ parameters — which, on a $1 \times 1$ spatial input, is just a linear layer wearing a convolution's clothes: each of the 10 outputs is a weighted sum of the 32 pooled channel values. Those 10 raw scores (logits) pass through log-softmax, which turns them into log-probabilities — ten numbers that sum, after exponentiating, to 1. The predicted digit is simply the index of the largest one. Log-softmax instead of plain softmax is a numerical-stability choice that also pairs cleanly with the negative-log-likelihood loss used during training; it changes nothing about which class wins, only how the winning margin is represented.

A 99.2%-accurate model is a black box unless you can ask it why. Grad-CAM (Gradient-weighted Class Activation Mapping) is the standard answer for convolutional networks: it turns the model's internal feature maps into a heatmap over the input image, showing which pixels the model's decision actually depended on. This model ships Grad-CAM as its one interpretability capability — the workspace runs it live on any digit you draw or upload.

Grad-CAM needs two things: a feature-map volume to explain, and a gradient telling it which channels in that volume mattered for the predicted class. For this model both come from conv_block4's output — the $8 \times 8 \times 32$ volume from the Architecture tab, the last spatial feature map before global average pooling discards position entirely. Picking the last convolutional block is deliberate: it's the deepest representation that still has spatial structure, so its activations are the most semantically rich map that can still be drawn as a picture.

Here is the computation, in the same notation as the rest of this page. Let $A^k$ be the $k$-th channel of conv_block4's $8 \times 8 \times 32$ output (so $k = 1, \dots, 32$, each $A^k$ an $8 \times 8$ grid), and let $y^c$ be the log-softmax score for the predicted class $c$ (one of the 10 digits). Grad-CAM first asks: if channel $k$'s activations all increased slightly, how much would the score for class $c$ change? That's a gradient, averaged over the $8 \times 8$ grid into one importance weight per channel:

Then the heatmap is a weighted sum of the 32 channel maps, passed through a ReLU to keep only the positive evidence — pixels that pushed toward the predicted class, not away from it:

Why the heat lands on ink strokes. Each cell of that $8 \times 8$ grid corresponds to a patch of the original 28×28 image — by conv_block4, after four 3×3 convolutions and one 2×2 pool, each cell's receptive field covers a sizeable neighbourhood, large enough to contain a stroke segment or a junction between strokes. The channels with the largest $\alpha_k^c$ are the ones whose detectors — built from the edge- and stroke-shaped kernels learned in the early blocks (the same vocabulary described in the convolution feature-hierarchy) — fired most strongly and most usefully for distinguishing this digit from the other nine. For a clean, centred digit, that combination reliably picks out the strokes themselves: a '7' lights up along its horizontal bar and diagonal, a '0' lights up as a ring along its loop. Background pixels — uniform, featureless — produce near-zero activation in every channel, so $\sum_k \alpha_k^c A^k$ is small there regardless of the weights, and the ReLU clips any residual negative response to exactly zero.

This is also a sanity check, not just a pretty picture. If Grad-CAM on a correctly-classified '7' highlighted the corners of the image instead of the stroke, that would be evidence the model is exploiting some dataset artefact (a border effect, a watermark) rather than the digit's shape — the kind of failure that test-set accuracy alone can hide. Draw a digit with an unusual stroke style in the workspace and watch whether the heat still tracks your ink: that's the model's 99.2% made legible, one prediction at a time.