A CNN from scratch: watching a convolutional network see, in your browser

Most introductions to convolutional neural networks open with a block diagram: a few boxes, a few arrows, an answer at the end. The interesting part is exactly what it hides: what happens inside the boxes?

This article goes the other way round. The classifier below is already running in your browser. There is no server behind it and no TensorFlow.js or ONNX Runtime; the whole forward pass is a little over two hundred lines of TypeScript.note The source lives in this site's repo under lib/ml: tensor.ts, ops.ts and sequential.ts, with zero dependencies. Draw a digit, and then we take it apart layer by layer.

An image is a grid of numbers

To a model, a greyscale image is a two-dimensional array in which every cell holds a number between 0 and 1: 0 is blank paper, 1 is ink. The small picture in the middle above is that array: $28 \times 28 = 784$ numbers.

The obvious approach is to flatten those 784 numbers and feed them into a fully connected layer. It works, but it's wasteful: a fully connected layer has no idea which pixels are neighbours. Shift the same "7" two pixels to the right and it becomes an entirely different input that has to be learned again.

A CNN starts from two assumptions that hold for almost every image:

• Locality. Meaningful features such as edges, corners and stroke ends involve only a small patch of neighbouring pixels.
• Translation equivariance. A vertical edge is the same thing in the top-left corner as in the bottom-right, so the same parameters should detect it in both. Convolution guarantees that shifting the input one cell right shifts the output one cell right; an answer that does not change with position at all (translation invariance) comes from the pooling that follows, and only approximately.

Build those two assumptions into the model's structure and you get convolution.

Convolution: a small window sliding over the image

Take a small matrix of weights, called a kernel, usually 3×3. Lay it over the top-left corner of the image, multiply the nine overlapping pairs, and add them up to get one number. Slide one pixel to the right and do it again. Once you've covered the whole image, those numbers form a new image called a feature map.

As a formula:

The figure below slows this down. The input is a bright vertical bar (its bottom two rows are shifted one pixel to the left), and the kernel has −1 down its left column and +1 down its right. Press Step and watch each output cell being computed.

This kernel computes "right minus left". Over a uniform region that is zero; only where dark on the left meets bright on the right does it become a large positive number. It is a vertical edge detector, and it uses the same nine numbers wherever the edge is.note Strictly speaking, what deep-learning frameworks call convolution is cross-correlation: the kernel isn't flipped. Because the weights are learned, flipping makes no difference, and the name stuck.

In code it is four nested loops. Below is a simplified single-channel version; the one in lib/ml/ops.ts adds loops over the batch and the channels, and its two innermost loops are exactly these:

for (let oy = 0; oy < oH; oy++) {
  for (let ox = 0; ox < oW; ox++) {
    let sum = bias;
    for (let ky = 0; ky < kH; ky++) {
      for (let kx = 0; kx < kW; kx++) {
        const iy = oy * stride + ky - padding;
        const ix = ox * stride + kx - padding;
        if (iy < 0 || iy >= h || ix < 0 || ix >= w) continue; // zero padding
        sum += x[iy * w + ix] * kernel[ky * kW + kx];
      }
    }
    out[oy * oW + ox] = sum;
  }
}

Change the numbers, change the feature

Nine numbers can do more than you'd expect. The input below is the digit you just drew (the sample 7 if you drew nothing); edit the kernel and see.

Classical computer vision designed these kernels by hand: Sobel, Laplacian, Gaussian. The key move in a CNN is to stop designing them: treat the nine numbers as parameters, initialise them randomly, and let gradient descent find the sets most useful for the task.

ReLU and pooling

A convolution is usually followed by two very small operations.

ReLU sets negative values to zero: $\mathrm{ReLU}(x) = \max(0, x)$. Convolutions stacked on convolutions, however many, are still one linear operation that could be folded into a single larger convolution, so the depth would buy nothing. ReLU is the non-linearity that makes depth mean something.

Max pooling replaces each 2×2 block with its largest value, halving the width and the height. Two things follow: later layers have a quarter as many pixels to process, and a feature gives the same output wherever it falls inside its block, so the model is less sensitive to small shifts.

The whole network

The model here has two convolutional blocks and one fully connected layer:

That is 9,098 parameters in a weights file of about 66 KB, reaching 98.6% accuracy on the MNIST test set.note Training was run once in PyTorch and the weights exported as JSON. The TypeScript forward pass has unit tests that compare it with PyTorch's output to within $10^{-4}$.

So few parameters is convolution's doing. For comparison I trained networks with one fully connected hidden layer on the same MNIST data:note Three seeds per width, 15 epochs, best epoch on the test set; script and output are in docs/research/mlp-baseline/.

In TypeScript the model is an array:

export const MNIST_CNN: LayerSpec[] = [
  { type: "conv2d", name: "conv1", inC: 1, outC: 8, kernel: 3, padding: 1 },
  { type: "relu", name: "relu1" },
  { type: "maxpool", name: "pool1", size: 2 },
  { type: "conv2d", name: "conv2", inC: 8, outC: 16, kernel: 3, padding: 1 },
  { type: "relu", name: "relu2" },
  { type: "maxpool", name: "pool2", size: 2 },
  { type: "flatten", name: "flatten" },
  { type: "dense", name: "fc", inF: 784, outF: 10 },
  { type: "softmax", name: "softmax" },
];

Sequential.forward() differs from an ordinary inference library in one deliberate way: it returns the output of every layer, not only the final answer. The figures that follow all depend on it.

What the network sees

These are the feature maps of every layer as your digit passes through the network. Brighter means a stronger response at that position.

Draw a few different digits:

1. In conv1 you can still recognise the digit; each map emphasises strokes in a different direction.
2. By conv2 a single map is hard to read. It encodes "some combination of strokes appears somewhere".
3. The final 16 × 7 × 7 = 784 numbers are everything the fully connected layer has to vote with.

Which pixels actually matter

High probability doesn't mean the model has understood anything. The most direct way to find out which part of the image it relies on is to cover that part and see how far the confidence falls.

The figure below slides a 4×4 blank patch across the image two pixels at a time, 13 × 13 = 169 positions in all. It re-runs the network at each position (skipping any where the patch covers only blank paper) and records how far the predicted class's probability drops.

Where it fails

Play with it for a while and you'll find ways to break it:

• An extra stroke, such as a bar through the middle of a 7 or a line under a 1, fools it easily: the training data has almost none of those.
• Something that isn't a digit still gets a confident answer. Softmax outputs always sum to 1; the model has no way to say "I don't know". On a production line or a surveillance feed a model will sooner or later see something outside its training data, and then high confidence is not the same as being right.

(Drawing small, or only in a corner, is fine: preprocessing crops, scales and then centres the drawing by its centre of mass, the same way MNIST was prepared.)

What comes next

This model has just over nine thousand parameters and ten classes. Stack the same bricks (convolution, non-linearity, downsampling) deeper and wider, add residual connections, and you have the backbones that do object detection, pose estimation and segmentation on edge devices today. The principle is unchanged; only the scale differs.

This model was trained beforehand and then brought in. To watch training itself happen in the browser, № 004 trains a Transformer from scratch while you watch its attention matrix take shape.