Capstone: Train a CNN and a Tiny GPT From Scratch (Grad-CAM, Perplexity)

Two models, one video. The first looks at a photo and says what is in it. The second reads a pile of old text and then starts writing its own. Both are built with the same tools you already know, and by the end you will have trained both yourself.

• Image classifier
• Tiny language model
• Both trained by you

Last time we finished customer segmentation and a recommendation system. Today is capstones nine and ten. Remember the rule for this whole capstone series: we report only results we measured ourselves. Any number you see me work through for arithmetic is labelled as an example, and your own runs will give different numbers.

• Report only measured results
• Examples are labelled
• No architecture re-teaching

Here is the plan. For images we run a full pipeline: a baseline convolutional network, augmentation, transfer learning, diagnostics, a confusion matrix, and Grad-CAM. For language we train a small character-level transformer, watch perplexity fall, sample with temperature, and be honest about its limits. We finish by comparing what each project costs in compute.

• Image pipeline
• Language model
• Compute comparison

We need a standard public dataset so anyone can compare notes. CIFAR-ten has sixty thousand colour images, thirty-two by thirty-two pixels, in ten classes such as airplane, cat, and truck. Fifty thousand are for training and ten thousand are for testing. It is small enough to train on one modest GPU, yet hard enough to be interesting.

• 60,000 images, 10 classes
• 32 by 32 colour pixels
• 50,000 train, 10,000 test

Before modelling, carve out a validation set from the training data, for example forty-five thousand to train and five thousand to validate. The test set stays locked away, exactly as in lesson thirty-five. Then normalise each colour channel using the mean and standard deviation computed on the training images only. A classic mistake is peeking at test statistics, which is a quiet form of leakage.

• Validation from training data
• Test set stays locked
• Stats from train only

Now the baseline. Using what we learned in lessons fifty-one and fifty-two, stack three convolution blocks, each with convolution, batch norm, ReLU, and max pooling, then a small linear classifier. Channels might go from three to thirty-two, sixty-four, and one hundred twenty-eight. Keep it deliberately simple, because a baseline exists to be beaten.

• Three conv blocks
• Pool halves the size
• Linear head, 10 logits

Train with cross-entropy loss and the Adam optimiser, a batch size of one hundred twenty-eight, and a few epochs. Log training loss, validation loss, and validation accuracy every epoch. Chance on ten balanced classes is ten percent accuracy, so that is your floor. Write down whatever accuracy you actually measure on validation, because every later improvement is judged against it.

• Cross-entropy and Adam
• Log every epoch
• Chance level is 10 percent

Read the curves like lesson fifty taught us. If training loss keeps falling while validation loss turns upward, the network is memorising. With only fifty thousand images, a baseline CNN usually does this. The gap between the two curves is the symptom, and more varied data is one of the best cures.

• Train down, validation up
• The gap means memorising
• Pick the best epoch

Augmentation creates new training examples by applying label-preserving changes. For CIFAR, the standard pair is a random crop after padding four pixels, and a random horizontal flip. A flipped cat is still a cat. But think before you augment: flipping a handwritten digit or rotating a street sign upside down could change the meaning.

• Random crop with padding
• Random horizontal flip
• Keep the label true

Two practical rules. Apply augmentation only to the training set, never to validation or test, because those must measure the real data. And the augmentation is random each epoch, so the network almost never sees the exact same pixels twice. Expect training accuracy to rise more slowly, and the validation gap to shrink.

• Never augment validation
• Random each epoch
• Gap should shrink

Diagnostics also catch bugs. Here is a classic sanity check: try to overfit a single small batch of thirty-two images. If the loss does not go near zero, something is broken, perhaps the learning rate or the labels. Also confirm you call model dot eval before validation, so batch norm and dropout behave correctly, and that initial loss is near the natural log of ten, about two point three.

• Overfit one tiny batch
• Initial loss near 2.30
• Switch to eval mode

A bigger lever is transfer learning, from lesson fifty-two. Take a ResNet eighteen already trained on ImageNet, whose early layers detect edges, textures, and simple shapes that are useful everywhere. Replace the final layer with a new one that outputs ten classes. You are borrowing millions of images worth of learning for free.

• ResNet-18 from ImageNet
• Early layers: edges, textures
• Swap the final layer

Do it in two stages. First freeze the pretrained body and train only the new head for a couple of epochs, with a normal learning rate. Then unfreeze everything and fine-tune with a much smaller learning rate, perhaps ten to a hundred times smaller, so you nudge the features instead of wrecking them. Use a different learning rate for head and body if you like.

• Stage one: train the head
• Stage two: unfreeze, small LR
• Different rates per part

Three details trip people up. The pretrained network expects ImageNet normalisation, so use its mean and standard deviation, not your own. It was trained on larger images, so you will often upscale the thirty-two pixel inputs, or adjust the first layer and stride. And keep augmentation on. Then compare to the baseline on the same validation set and record the real difference you measure.

• Use ImageNet normalisation
• Mind the input size
• Compare on same validation set

A single accuracy number hides where the model fails. Now, and only now, evaluate once on the test set, and build a confusion matrix. Rows are true classes, columns are predicted classes. The diagonal holds correct predictions, and every off-diagonal cell is a specific kind of mistake.

• Rows: true class
• Columns: predicted class
• Off-diagonal: mistakes

Let us read one with example numbers. Suppose there are one thousand test cats. If eight hundred are labelled cat, one hundred are labelled dog, and one hundred are scattered elsewhere, then recall for cat is eight hundred over one thousand, which is zero point eight. Precision looks down the column instead: of everything predicted cat, how many truly were cats. These are lesson sixteen's metrics, now per class.

• Recall: read along the row
• Precision: read down the column
• Example: 800 / 1000 = 0.8

Do not stop at the matrix. Pull out the actual images from the worst cell, for example cats predicted as dogs, and look at them with your own eyes. You may find blurry pictures, odd crops, or even wrong labels in the dataset. Sorting mistakes by model confidence also helps: confident errors are the most informative ones.

• Open the worst cell
• Check for bad labels
• Sort by confidence

To see why the network decides, we use Grad-CAM. Take the last convolution layer, with feature maps A sub k. Compute the gradient of the class score with respect to each map, and average it over all positions to get a weight alpha k. Then form a weighted sum of the maps and apply ReLU, keeping only evidence that pushes the class score up.

• Gradients of class score
• Average to weights alpha k
• Weighted sum, then ReLU

The result is a coarse heat map, which you upsample and overlay on the picture. A healthy model lights up the animal. A worrying one lights up the background, such as grass behind a horse or water behind a ship. That is shortcut learning, and it tells you the model may be right for the wrong reason. Treat Grad-CAM as a clue, not proof.

• Heat map over the image
• Check object versus background
• A clue, not proof

Now the second project. We will train a character-level GPT on a small public-domain text, such as the tiny Shakespeare file of roughly one million characters, or any Project Gutenberg book. Character-level means the vocabulary is tiny, about sixty-five symbols for that file, so we skip the tokenizer machinery from lesson fifty-four. A small byte-pair encoding version works too, with a bigger vocabulary and shorter sequences.

• Public-domain text
• About 65 characters
• BPE is an alternative

The task is next-token prediction. Take a window of, say, two hundred fifty-six characters as input. The target is the same window shifted one position to the right. So every position in the window is a training example, and a single batch teaches the model many predictions at once. Split the text by position, the first ninety percent for training and the last ten percent for validation, never by shuffling windows.

• Target is input shifted by one
• Every position is an example
• Split by position, not shuffle

The model uses the transformer block from lesson fifty-six, and we will not re-teach it. Stack token embeddings plus position embeddings, then say six blocks of masked self-attention and feed-forward layers with residuals and layer norm, then a final linear layer to vocabulary logits. A reasonable small setup is embedding size three hundred eighty-four, six heads, and six layers, around ten million parameters. Pick sizes your hardware can handle.

• Embeddings plus positions
• Six masked blocks
• Linear head to logits

The loss is cross-entropy between predicted next-character distribution and the true character. Here is a free sanity check. A model that guesses uniformly over sixty-five characters has loss equal to the natural log of sixty-five, about four point one seven. If your first step prints something far from that, your initialisation or labels are wrong.

• Uniform guess over 65
• Loss = ln 65 = 4.17
• Check step zero

Perplexity is just the exponential of the cross-entropy loss, measured in nats. It reads as the effective number of choices the model is torn between. Uniform guessing gives perplexity sixty-five. As an example, a validation loss of one point five gives e to the one point five, about four point four eight, meaning the model is as unsure as choosing among roughly four and a half characters.

• Perplexity = exp(loss)
• Uniform guess: 65
• Example: exp(1.5) = 4.48

The training loop looks like lesson fifty. Use AdamW, a short warmup, a cosine decay of the learning rate, gradient clipping at one, and evaluate on validation every few hundred steps. Plot training and validation loss, plus perplexity. On a small file you will see the same story as before: validation eventually stops improving while training keeps dropping. Save the checkpoint with the best validation loss.

• AdamW, warmup, cosine decay
• Clip gradients at 1.0
• Keep best validation checkpoint

Common mistakes here are subtle. Forgetting the causal mask lets the model peek at the answer, and the loss collapses suspiciously fast. Using overlapping windows between train and validation leaks text across the boundary. And if loss becomes not a number, lower the learning rate, check clipping, and consider mixed precision settings. Always debug on a tiny slice first.

• Missing causal mask
• Leaking across the split
• Exploding loss

Generation works by repeating one step. Feed in a prompt, get logits for the next character, turn them into probabilities, sample one character, append it, and repeat. The sampling knob is temperature, which divides the logits before the softmax. Low temperature sharpens the distribution, and high temperature flattens it.

• Predict, sample, append
• Repeat character by character
• Temperature divides logits

Here is the arithmetic. Take three logits: two, one, and zero. At temperature one, the probabilities are about zero point six seven, zero point two five, and zero point zero nine. At temperature zero point five the logits become four, two, zero, and the probabilities sharpen to about zero point eight seven, zero point one two, and zero point zero two. At temperature two they flatten to roughly zero point five one, zero point three one, and zero point one nine.

• T = 0.5: 0.87, 0.12, 0.02
• T = 1: 0.67, 0.25, 0.09
• T = 2: 0.51, 0.31, 0.19

What does that do to text? Very low temperature gives safe but repetitive output, often looping the same phrase. Temperature around one gives the most natural variety. High temperature produces creative nonsense and invented words. Generate samples at several temperatures from the same prompt, and quote them exactly as they come out.

• Low: repetitive
• Around 1: varied
• High: garbled

Now the honest part. A small model trained on one book learns spelling, punctuation, the layout of a script, and local word patterns. It does not learn meaning, plot, or facts. Expect text that looks like the source in shape but wanders in sense, with invented words and characters that lose track of who is speaking. That is not a bug in your code; it is what a few million parameters and one million characters can do.

• Learns form, not meaning
• Loses long-range coherence
• Few parameters, little data

A related risk is memorisation. With a tiny dataset, a model can reproduce long passages of its training text word for word, and the validation curve is how you notice. Low validation loss with fluent output may simply mean copying. Compare generated passages against the source before you claim the model wrote anything original.

• Small data invites copying
• Compare output to the source
• Validation loss is your guard

Let us compare the two projects. Training cost is roughly six times the number of parameters times the number of tokens processed. As an illustrative example, a ten-million-parameter model reading one million characters in an epoch needs about six times ten to the seventh times ten to the sixth, which is six times ten to the thirteenth floating point operations per epoch. That is small on a modern GPU, though real speed depends on your hardware and code.

• FLOPs about 6 x N x D
• Example: 6 x 10^13 per epoch
• Hardware sets real time

The image project has a different shape. A pretrained ResNet-eighteen has about eleven million parameters, but each image is pushed through many convolutions at several resolutions, and you repeat for fifty thousand images over many epochs. Fine-tuning is cheaper than training from scratch because you need fewer epochs. The language model has the quadratic cost of attention in context length, so doubling the context roughly quadruples that part.

• CNN: images times epochs
• Fine-tuning needs fewer epochs
• Attention: quadratic in context

Do not trust my estimates or anyone else's. Measure your own: wall-clock time per epoch, peak GPU memory, and the number of parameters, along with the final metric. Put them in a table for each project, and record the seed, hardware, and library versions. That table is your honest result, and it lets anyone reproduce or challenge your claims.

• Time, memory, parameters
• Record seed and hardware
• Report only your numbers

One last lesson from putting both projects side by side. The image task rewarded borrowing: a pretrained body beat starting from scratch on little data. The language task showed the opposite pressure: with only a book of text, scale of data and model is the limiting factor, which is why lesson fifty-seven talked about pretraining. Knowing which bottleneck you face decides what to try next.

• Images: borrow features
• Language: data and scale
• Diagnose before you fix

Let us recap. For images, we split carefully, trained a baseline, used augmentation and transfer learning, read diagnostics, built a confusion matrix, and used Grad-CAM to check what the model looked at. For language, we trained a character-level GPT, tracked loss and perplexity, sampled with temperature, and stayed honest about the limits.

• Image pipeline end to end
• GPT, perplexity, temperature
• Compare compute, report honestly

Your homework is to run both projects and fill in your own table. Change one thing at a time, such as the augmentation, the temperature, or the context length, and write down what happened. Next lesson, capstones eleven and twelve: time-series forecasting, and an end-to-end deployed machine learning app.

• Run both projects yourself
• Change one thing at a time
• Next: capstones 11 and 12