Training a CNN
in PyTorch

OVERVIEW

What you'll learn

In Part 1 we built the pipeline to prepare the datasets - A custom Dataset, transforms, and a DataLoader that hands out augmented, shuffled batches of Oxford 102 Flowers with 8K images images across 102 classes. We also split the dataset into training, test and validation and show how the transforms were applied differently for training and validation datasets. If you haven't read it, you may want to read that first.

In this part, we will build a convolutional neural network (CNN) from scratch in PyTorch and train it on the Oxford 102 Flowers dataset. We will also inspect and debug the model and learn how to improve CNNs.

WHY CNNS?

Why a stack of linear layers can't see a flower?

Linear layers treat every pixel as independent, CNNs do not.

To train an image with 224 x 224 x 3 channels, a stack of nn.Linear layers would require 224 * 224 * 3 = 150,528 input nodes. This is computationally expensive and leads to overfitting due to the large number of parameters.

And any hidden layer added increases the complexity and parameters really quickly. If you were to add even one hidden layer with 2048 nodes, the number of parameters would explode. This is already 300M+ parameters

Additonally, Linear layers treat every pixel as independent. The model sees thousands of separate numbers with no understanding that neighbouring pixels form petals, edges, and veins.

To contrast, how do you look at image? You don't identify a flower pixel by pixel either. You notice shape, contrast, texture. That's what CNNs are made for. Understand pattern from the images

Let's understand how CNNs do it, via convolution.

CONVOLUTION

The C in CNNs

Convolution is a filter

At each position, multiply the filter values with the pixel values underneath and add them together. Slide it across the image and every pixel gets a new value based on its neighbours. That process is a convolution.

Different weights highlight different patterns. Take a vertical edge filter. Similar pixels left and right cancel out. A sharp contrast — dark left, bright right — produces a strong output. Rotate the same idea and you get a horizontal edge detector.

The output isn't a picture. It's an array of numbers showing how strongly that filter reacted to each part of the image — a feature map.

Why the network learns its own filters

You could design filters by hand. But which filters separate a pink primrose from a wild pansy? You don't know, and you shouldn't have to. The model learns the filters and tunes them to the patterns that distinguish your classes.

The five arguments to nn.Conv2d

nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1)

• in_channels — colour channels coming in. Three for RGB.
• out_channels — how many filters this layer learns. Each detects a different feature.
• kernel_size — the size of each filter. 3×3 sees every pixel with its immediate neighbours.
• stride — how far the filter moves each step. 1 checks every pixel; 2 halves the output.
• padding — adds zeros outside the border so the filter can centre on an edge pixel.

What a filter actually does to a picture

That is the mechanic. Here is what it looks like on a real photo. The same image runs through two filters that differ only in how their nine numbers are arranged — and each one keeps a different set of edges.

horizontal edges

vertical edges

Both filters are the same nine numbers, transposed. The horizontal kernel lights up the rooftop deck, the ledges, and the waterline; the vertical one lights up the towers' sides. The kernel decides which edges survive.

Making the Image smaller

MaxPooling

# halves H and W, no weights
nn.MaxPool2d(kernel_size=2)  

Max pooling (MaxPool2d) takes a 2×2 area, keeps the largest value, discards the rest. The output is a quarter of the size. Repeat across every 2×2 group and you end up with a smaller output. This is safe because your filters have already extracted the important features from the original image, so pooling just compresses each filtered image, keeping the most significant information.

The highlighted window — 168, 0, 0, 64 — keeps 168 and drops the other three. Repeat for all four windows and the 4×4 map becomes 2×2.

CHAINING CONV2D, RELU and MAXPOOL2D

The First Architecture

Four blocks of convolution, ReLU, pool. Then three linear layers. Textbook.

model = nn.Sequential(
    nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),

    nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),

    nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),

    nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),

    nn.Flatten(),
    nn.Linear(128*196, 2048),
    nn.Dropout(0.5),
    nn.Linear(2048, 512),
    nn.Dropout(0.5),
    nn.Linear(512, 102),
)

Everything the convolutions learn is squeezed through one 25,088-wide vector. Four convolution layers hold 0.2% of the model.

Channels grow 3 → 16 → 32 → 64 → 128. Padding 1 with kernel 3 keeps the size, so only the pooling layers shrink anything. Before it can train, two things need setting up: the defences that stop it memorising, and the loss and optimizer that drive it.

STOPPING THE MODEL FROM MEMORISING

Regularization

Dropout, and the husky that got called a wolf

Ribeiro et al. (2016): a husky classified as a wolf, not because it looked like one, but because the model had learned that snow means wolf. Some neurons become snow detectors and the rest lean on them. That's co-adaptation — the model leaning on a shortcut instead of learning body shape.

Look at almost any set of wolf photos and the background gives the game away: snow, again and again. A model can score well by learning "white background" instead of "wolf" — right for the wrong reason, until the first wolf on grass.

We avoid this co-adaption with Dropout Layers. During training, a dropout layer randomly deactivates a fraction of the neurons — around half. It sounds destructive. But the advantage is that, Dropout makes shortcuts risky. If the snow detector might vanish on any given batch, the network has to find other cues — the ones that actually matter. The technique is from Srivastava et al. (2014).

Each training step, dropout removes a random subset of units and every connection into or out of them. The network can't build a fixed dependency between two specific neurons if either might be gone next batch, so it spreads the work across many. At test time all units are back.

In practice dropout rates run 0.2 to 0.5, placed after the activation and before the final classification layer

When to not confuse with Codaptation

If every wolf image has snow and no dog image does, that's a dataset problem, not overfitting. The model is learning the only pattern it can see. Dropout helps when patterns are mixed, not when they're absent.

LOSS AND OPTIMIZER

Loss and Optimizer

The loss is CrossEntropyLoss — standard for multi-class classification. The optimizer is Adam.

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)

loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0005, weight_decay=0.0005)

CrossEntropyLoss expects raw scores, not probabilities — which is why the model ends at a bare linear layer with no softmax. Adding one here would be a bug that trains quietly and badly

Weight decay

weight_decay=0.0005 discourages the network from using very large weights. Large weights are often a sign the model is memorising specific training examples rather than learning features that generalise, so a small penalty for large weights nudges it toward simpler, more robust solutions.

RUNNING THE LOOP

Training

model.train() and model.eval() switch dropout and batch norm between modes — leave dropout on during validation and the accuracy is wrong. And optimizer.zero_grad() comes first, because gradients accumulate by default.

def train_model():
    model.train()
    train_loss = correct = total = 0
    for batch, (imgs, lbls) in enumerate(train_loader):
        imgs, lbls = imgs.to(device), lbls.to(device)
        optimizer.zero_grad()
        preds = model(imgs)
        loss = loss_function(preds, lbls)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = preds.max(1)
        total += lbls.size(0)
        correct += predicted.eq(lbls).sum().item()

        print(f'BATCH {batch+1} | TRAIN LOSS: {train_loss/(batch+1):.3f} | ACCURACY: {correct*100/total:.2f} %')

    torch.save({"epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict()},
               ckpt_dir / f"epoch_{epoch:03d}.pt")


def eval_model():
    model.eval()
    val_loss = correct = total = 0
    with torch.no_grad():
        for batch, (imgs, lbls) in enumerate(val_loader):
            imgs, lbls = imgs.to(device), lbls.to(device)
            preds = model(imgs)
            loss = loss_function(preds, lbls)
            val_loss += loss.item()
            _, predicted = preds.max(1)
            total += lbls.size(0)
            correct += predicted.eq(lbls).sum().item()

    return val_loss/BATCH_SIZE, correct/total

for epoch in range(30):
    train_model()
    vl, acc = eval_model()
    print(f'EPOCH: {epoch+1} | VAL LOSS:{vl:.3f} | VAL ACCURACY: {acc*100:.2f}%')
    report = run_diagnostics(model, val_loader, train_dataset, val_dataset, device)
    train_vs_val_gap(model, train_dataset, val_loader, device)

The diagnostics the loop calls

Those last two lines are where every per-class number in this post comes from. Both run on the validation set after each epoch. Trimmed to the parts this post quotes:

@torch.no_grad()
def collect_predictions(model, loader, device):
    model.eval()
    ys, ps, top5s, confs = [], [], [], []
    for imgs, lbls in loader:
        logits = model(imgs.to(device))
        ys.append(lbls)
        ps.append(logits.argmax(1).cpu())
        top5s.append(logits.topk(5, dim=1).indices.cpu())
        confs.append(logits.softmax(1).max(1).values.cpu())
    return [torch.cat(x).numpy() for x in (ys, ps, top5s, confs)]


def run_diagnostics(model, val_loader, train_dataset, device, n_classes=102, show=15):
    y, pred, top5, conf = collect_predictions(model, val_loader, device)
    correct = (pred == y)
    top5_correct = (top5 == y[:, None]).any(1)
    train_counts = np.bincount(labels_of(train_dataset), minlength=n_classes)

    print(f'  top-1 accuracy  : {correct.mean()*100:.2f}%')
    print(f'  top-5 accuracy  : {top5_correct.mean()*100:.2f}%')
    print(f'  mean confidence : {conf.mean():.3f} '
          f'(correct {conf[correct].mean():.3f} / wrong {conf[~correct].mean():.3f})')

    # collapse check — is it answering with only a handful of classes?
    pred_counts = np.bincount(pred, minlength=n_classes)
    used = (pred_counts > 0).sum()
    print(f'  distinct classes predicted : {used}/{n_classes}')
    if used < n_classes * 0.7:
        for c in np.argsort(-pred_counts)[:5]:
            print(f'    {NAMES[c]:<28} {pred_counts[c]*100/len(y):5.1f}% of all predictions')

    # per-class accuracy, sorted worst to best
    rows = [(c, int((y == c).sum()), correct[y == c].mean(),
             top5_correct[y == c].mean(), int(train_counts[c]))
            for c in range(n_classes) if (y == c).any()]
    rows.sort(key=lambda r: (r[2], r[1]))
    for title, block in [('WORST', rows[:show]), ('BEST', rows[-show:][::-1])]:
        print(f'\n{title} {show} CLASSES')
        for c, n, a1, a5, tn in block:
            print(f'  {NAMES[c]:<28} {n:>5} {a1*100:>6.1f}% {a5*100:>6.1f}% {tn:>7}')

    # is accuracy just tracking how much data each class had?
    tn = np.array([r[4] for r in rows], dtype=float)
    a1 = np.array([r[2] for r in rows], dtype=float)
    print(f'\n  corr(train images, val accuracy) = {np.corrcoef(tn, a1)[0,1]:+.3f}')

def train_vs_val_gap(model, train_dataset, val_loader, device, max_batches=10):
    model.eval()
    tr = accuracy(DataLoader(train_dataset, batch_size=128), limit=max_batches)
    va = accuracy(val_loader, limit=10**9)
    print(f'  train accuracy (in eval mode): {tr*100:.2f}%')
    print(f'  gap                          : {(tr-va)*100:+.2f} pts')
    if tr - va > 0.15:
        print('  -> OVERFITTING. Add augmentation, weight decay, early stopping.')
    elif tr < 0.75:
        print('  -> UNDERFITTING. The model is the ceiling, not the data.')
    return tr, va

The train side samples ten batches, not the whole training set, so the gap is an estimate. That is deliberate — a full pass over 5,737 images every epoch would cost more than the training step it is diagnosing.

The first architecture, over thirty epochs

Run that loop on the first architecture and here is what comes back.

V1 over 30 epochs — hover for the numbers. Training accuracy climbs past 81%. Validation accuracy crosses 55% at epoch 13 and then goes sideways, while validation loss bottoms out at epoch 18 and starts climbing back.

Over the last ten epochs validation accuracy moves between 58.14% and 61.81%. Best is 61.81% at epoch 28. Training loss keeps falling the whole time, 0.887 down to 0.623. Loss improving while validation accuracy does not is the signature of a model spending its capacity on the training set.

The problem is the single number 128*196.

Reading shapes, and the 51-million-parameter layer hiding in one line

Debugging

Where does 196 come from, and what does it cost?

Step 1: Debug Matrix Shapes

Images arrive at 224×224 from CenterCrop(224). Four pooling layers, each halving: 224 → 112 → 56 → 28 → 14. The last convolution outputs 128 channels. So the flattened vector is 128 × 14 × 14, and 14 × 14 is 196.

If the shapes don't match you, will likely run into errors like this:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (128x25088 and 2048x2048)

The first pair is what arrived — batch of 128, flattened to 25,088. The second is what the layer expected. Two numbers, and the fix is whichever one you can defend.

img, label = next(iter(train_dataset))
x = img.unsqueeze(0)

for layer in model:
    x = layer(x)
    print(f'{layer.__class__.__name__:16} {tuple(x.shape)}')

This works because PyTorch builds the graph as the code runs — more on that in a moment.

Step 2: Counting Model Parameters

nn.Linear(25088, 2048) is 51,380,224 weights. The whole model is 52,581,126. One line holds 97.7% of the network.

Trainable parameters: 52,581,126

  all four conv layers        97,440      0.2%
  Linear(25088, 2048)     51,382,272     97.7%
  Linear(2048, 512)        1,049,088      2.0%
  Linear(512, 102)            52,326      0.1%

Version 2: nine times smaller, thirteen points better

The Rebuild

Same data, same optimizer, same 30 epochs. Different shape of network.

model = nn.Sequential(
    nn.Conv2d(3, 32, kernel_size=3), nn.ReLU(), nn.BatchNorm2d(32),
    nn.Conv2d(32, 32, kernel_size=3), nn.ReLU(), nn.BatchNorm2d(32),
    nn.Conv2d(32, 32, kernel_size=5, stride=2, padding=2), nn.ReLU(), nn.BatchNorm2d(32),
    nn.Dropout(0.4),

    nn.Conv2d(32, 64, kernel_size=3), nn.ReLU(), nn.BatchNorm2d(64),
    nn.Conv2d(64, 64, kernel_size=3), nn.ReLU(), nn.BatchNorm2d(64),
    nn.Conv2d(64, 64, kernel_size=5, stride=2, padding=2), nn.ReLU(), nn.BatchNorm2d(64),
    nn.Dropout(0.4),

    nn.Conv2d(64, 128, kernel_size=3), nn.ReLU(), nn.BatchNorm2d(128),
    nn.Conv2d(128, 128, kernel_size=3), nn.ReLU(), nn.BatchNorm2d(128),
    nn.Conv2d(128, 128, kernel_size=5, stride=2, padding=2), nn.ReLU(), nn.BatchNorm2d(128),
    nn.Dropout(0.4),

    nn.Conv2d(128, 256, kernel_size=5, stride=2, padding=2), nn.ReLU(), nn.BatchNorm2d(256),
    nn.Conv2d(256, 512, kernel_size=5, stride=2, padding=2), nn.ReLU(), nn.BatchNorm2d(512),

    nn.AdaptiveAvgPool2d(4),
    nn.Flatten(),
    nn.Dropout(0.4),
    nn.Linear(8192, 102),
)

trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Trainable parameters: {trainable_params:,}")

Trainable parameters: 5,760,166

Same picture, redrawn. No single layer dominates: stages 1–3 hold 14.3% of the parameters, the 512-channel convolution 56.9%, the classifier 14.5%. The largest block is now a convolution that actually learns features.

Five changes, and what each one is for

• AdaptiveAvgPool2d(4) — averages whatever grid arrives down to 4×4. The flatten is always 512 × 4 × 4 = 8,192, whatever the input size. One head, 835K parameters instead of 51M, and you never compute that number by hand again.
• BatchNorm2d after every convolution — rescales each channel to zero mean and unit variance across the batch, which keeps activations from drifting as depth grows. Standard practice for a stack this deep; I did not run V2 without it.
• stride=2 convolutions replace MaxPool2d — the network learns how to downsample instead of always keeping the maximum.
• two 3×3 convolutions before each downsample — more depth per stage, so features compose before the picture shrinks.
• Dropout(0.4) between stages, not only in the classifier — regularisation applied where the features are, not just at the end.

All five changes shipped together, so this comparison does not tell you which one earned the thirteen points. Treat the table as the outcome of a rewrite, not as evidence for any single line in it.

Note the 3×3 convolutions have no padding, so each one shaves two pixels off every side. Follow 192 through the stack: 192 → 190 → 188 → 94, then 92 → 90 → 45, then 43 → 41 → 21, then 11, then 6. AdaptiveAvgPool2d takes that 6×6 to 4×4 and the arithmetic stops mattering.

THE REBUILT MODEL, MEASURED

After the Rebuild

Two accuracy numbers appear from here on. Top-1 is the ordinary one: the highest-scoring class is the right one. Top-5 asks whether the right class is anywhere in the model's five best guesses — a useful second reading, because a model that is close but not first is a different problem from one that is lost.

The loss curve

Training loss falls from 3.741 to 0.188 and keeps falling. Validation loss flattens around epoch 20 near 1.1. When one keeps dropping and the other doesn't, you're past the useful part of training.

The accuracy curve

Validation top-1 climbs from 9.87% to a best of 74.80% at epoch 28 and top-5 reaches 91.52%. The shaded band is the gap between train and validation — the thing to watch.

What the epochs actually say

Epoch 1 — collapse. The model predicted only 57 of 102 classes, and 28% of every prediction it made was "rose". The train/validation gap was +1.8 points, which reads as underfitting: the model was the ceiling, not the data.

Epochs 2 to 7 — the collapse resolves. Distinct classes predicted goes 57 → 86 → 98 → 102. Once the network stops hedging on a few common flowers, accuracy triples.

Epochs 9 to 15 — the crossover. The gap passes 10 points at epoch 9 and 20 points by epoch 15. Training accuracy is now pulling away. The model has stopped underfitting and started memorising.

Epochs 20 to 30 — diminishing returns. Measured to epoch 30, training gains 5.4 points and validation gains 1.5, and the gap widens to 25.3. The best validation accuracy — 74.80% — lands at epoch 28, and it is 3.2 points above where the run finishes. Your best model is not your last one, which is why the loop checkpoints every epoch.

V1 against V2, epoch for epoch

V2 passes V1's final 60.02% at epoch 15 and keeps climbing for another thirteen epochs. Nine times fewer parameters, thirteen points better.

What thirty epochs costs on a T4

Both runs are Kaggle notebooks on a single T4. The V2 log gives per-epoch wall-clock times, and they're remarkably flat

Which flowers it actually learned

Press play to run the leaderboard through all thirty epochs. By epoch 30 ten classes are perfect; hover any row for its top-5 score and how many images it had to learn from — wallflower gets all 29 of its validation images right, off 138 training images.

Sixty-one different classes pass through the top fifteen over thirty epochs, with an average of eight swapping in or out each epoch. Model training is fairly stable until about epoch 20.

Class size explains some of it. Split the 102 classes into quartiles by training-set size: the smallest quartile, 28 to 35 images, averages 60.6%; the largest, 65 to 180 images, averages 77.8%. A seventeen-point spread is real. But the correlation across all classes is only +0.218 — weak. If imbalance were the bottleneck you would expect something north of +0.4. So more data per class would help, and it would not be the fix.

Similarity explains the rest. The confusion pairs at epoch 30 are not random: primula read as wallflower six times, bougainvillea as petunia four, sunflower as barbeton daisy three, colt's foot as marigold three. Those are yellow-centred composites and open trumpets — flowers that genuinely resemble each other at 192 pixels. The dataset's own authors note it contains several very similar categories. That is the part a bigger model, not a bigger dataset, has to solve.

What to check before you start a thirty-epoch run

Inspection

print(model) shows layer names, types, and settings — good for spotting structural mistakes. It won't tell you how big the model is or where the weight sits.

# does one image survive the whole stack?
img, label = next(iter(train_dataset))
print(model(img.unsqueeze(0)).shape)          # torch.Size([1, 102])

# how big is it, and where does the weight live?
total = sum(p.numel() for p in model.parameters())
print(f'total parameters: {total:,}')

for name, p in model.named_parameters():
    print(f'{name:24} {tuple(p.shape)}  {p.numel():,}')

Three things worth knowing about that output. model.parameters() returns a generator, so nothing loads into memory until you iterate. numel() gives the element count per tensor. And a linear layer's shape reads outputs first — (102, 8192), one row per output neuron.

For nested models, children() shows only the top level while modules() walks the whole tree. Folders versus everything inside them.

Why you can print a tensor in the middle of a model at all

Dynamic Graphs

When you chain layers together, you are really writing one big equation. Conv2d is thousands of multiplications and additions. ReLU zeroes out the negative values. Written out one step at a time — multiply, add, zero, multiply, add the bias — that step-by-step breakdown is called a computation graph.

The framework records every operation for a reason. To train the model, it has to walk backwards through the graph and use the chain rule to work out how to adjust each parameter.

Traditional Libraries

Older frameworks asked you to define the entire graph up front, before any data moved through it:

• You wrote out every operation and every connection first.
• Once defined, the structure was locked.
• The upside: a framework that knows the full graph ahead of time can optimise it for speed and memory.
• The downside: no flexibility. You are describing a fixed equation, so there are no loops, no conditionals, and no way to print a value from the middle.

The PyTorch way: build the graph as the code runs

PyTorch builds the graph while your code runs. You can put branching logic straight inside the forward method, and the if statement does not just choose the logic — it shapes the graph itself.

def forward(self, x):
    x = self.features(x)

    # the branch taken decides the graph that gets built
    if x.mean() > self.threshold:
        x = self.detail_head(x)      # tricky sample, heavier path
    else:
        x = self.simple_head(x)      # easy sample, cheaper path

    return x

Here is what happens each time forward runs:

• PyTorch records exactly what happened — every multiplication, addition, layer, and branch.
• That graph is used for backpropagation, and the parameters get updated.
• Then the graph is thrown away.
• The next batch builds a brand new graph, even if it takes a different path.

What this makes possible

• Variable input shapes. A three-word sentence and a fifty-word sentence can go through the same model.
• Easy debugging. No special debug mode — it is just Python, so you add a print.
• Input-adaptive models. A cheaper path for easy inputs and a heavier path for hard ones.

There is a small performance cost to building the graph each time, but for most work the flexibility is worth it. Static frameworks make you think like a compiler; PyTorch lets you think like a Python programmer.

Where nn.Sequential fits

nn.Sequential gives that flexibility back up. It is a fixed, straight-line path: no conditionals, no loops, and no returning two outputs. Both models in this post are Sequential, and that is the right choice while the network is a simple pipeline.

When you need branching, you move to nn.Module: __init__ holds the layers, and forward holds the flow.