A convolutional network that separates cats from dogs, or good parts from defective ones on a production line, sounds like a serious project. It is not. A working image classifier is roughly eighty lines of PyTorch, and most of those lines are patterns you will recognise after the second read.

What follows is a complete build in about thirty minutes of typing. The dataset is CIFAR-10, which downloads through torchvision and hands you 60,000 colour images across ten classes — enough to watch real learning happen without waiting overnight. The same code runs on any folder of images sorted into class subfolders, so when you finish here, you can point it at your own pictures.

What You Need Before You Start

Python 3.9 or newer, PyTorch and torchvision installed, and a terminal you are comfortable in. A GPU helps but is not required: this network trains on a laptop CPU in a few minutes per run. If PyTorch is not installed yet, use the selector on the official installation page and copy the command it generates rather than guessing at CUDA versions — mismatched builds are the single most common reason a first tutorial fails.

Then confirm your imports work: import torch, torchvision, torch.nn as nn, torch.optim as optim, plus from torchvision import datasets, transforms and from torch.utils.data import DataLoader. If that prints nothing, you are ready.

Step 1 — Load the Data and Look at It

Every image must reach the network as a tensor of consistent size with comparable pixel statistics. torchvision handles that with a transform pipeline: transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))]). ToTensor scales pixels to 0–1 and reorders dimensions to channels-first; Normalize shifts them to roughly −1 to 1.

Wrap CIFAR-10 in datasets.CIFAR10(root="./data", train=True, download=True, transform=...) and hand it to a DataLoader with batch_size=64 and shuffle=True for the training set. Before you write a single layer, pull one batch and check its shape. You want torch.Size([64, 3, 32, 32]): batch, channels, height, width. Then plot four images alongside their labels. Two minutes of looking catches wrong paths, wrong label mappings, and corrupted downloads before they turn into confusing loss curves.

Step 2 — Define a Small Convolutional Network

Three convolutional blocks are plenty for 32×32 images. Each block is a convolution with padding=1 to preserve spatial size, a ReLU, and a 2×2 max pool that halves width and height. So: 3 channels to 32, then 32 to 64, then 64 to 128. After three pools, a 32×32 image is 4×4 with 128 channels, which flattens to 2048 numbers. Two linear layers finish the job — 2048 to 128, then 128 to 10 classes.

That is under 300,000 parameters. It will not win a benchmark, but it reaches respectable accuracy on CIFAR-10 and trains fast enough that you can run experiments all afternoon.

Two details that matter

Keep the final linear layer free of softmax. nn.CrossEntropyLoss expects raw logits and applies the log-softmax internally; adding your own softmax first quietly degrades training, and it is one of the most common early mistakes. Also count your parameters with sum(p.numel() for p in model.parameters()). Knowing your budget helps you tell "too small" apart from "overfitting" later.

Step 3 — Write the Training Loop

The inner loop is five lines you will type a thousand times:

1. optimizer.zero_grad() to clear gradients from the previous step.
2. outputs = model(images) for the forward pass.
3. loss = criterion(outputs, labels) to score the predictions.
4. loss.backward() to compute gradients.
5. optimizer.step() to update the weights.

Surround that with a loop over the DataLoader inside a loop over epochs. Ten to fifteen epochs with optim.Adam(model.parameters(), lr=1e-3) is a sensible start. Print the running loss every hundred batches; on CIFAR-10 it should fall from around 2.3 toward 0.5 or below within a handful of epochs. If it sits flat, your learning rate is likely wrong, or your labels are not lining up with your inputs.

Step 4 — Measure Accuracy Honestly

A loss number tells you the model is improving. Accuracy tells you whether it is useful. Keep a held-out set — CIFAR-10's own test split works — and after each epoch switch the model to evaluation mode with model.eval(), wrap the pass in torch.no_grad(), and divide correct predictions by the total.

Then compare the two numbers. If training accuracy is 98% and test accuracy is 55%, the model has memorised the training set. The fixes, in order of effort: more data, random augmentation such as flips and small crops, dropout before the linear layers, or simply fewer epochs. If both numbers are low, the model is underfitting — give it more capacity or train longer.

Step 5 — Classify a Single New Image

Training ends with a checkpoint: torch.save(model.state_dict(), "cifar_net.pt"). Reload it later with the same architecture and model.load_state_dict(...).

To predict one image, open it with PIL, apply the exact transforms used during training, and add a batch dimension with img.unsqueeze(0). Run the forward pass under no_grad, take torch.argmax over the output, and map the index back to a class name. If you want a confidence score, wrap the logits in torch.softmax and read the largest value — but treat anything above 0.9 as a hint rather than a guarantee, especially on images that look nothing like your training data.

Mistakes That Cost Beginners an Evening

• Forgetting optimizer.zero_grad(), which accumulates gradients and makes the loss jump around.
• Never calling model.eval(), so dropout and batch normalisation keep behaving as they do in training and your test accuracy looks worse than it is.
• Using a different resize or normalisation at inference than at training time.
• Passing labels as floats instead of longs to CrossEntropyLoss.
• Shuffling the test set and then trying to compare predictions to labels.
• Training for 200 epochs on a small dataset and calling the resulting memorisation "accuracy".

Where to Go Next

Swap CIFAR-10 for your own images by pointing datasets.ImageFolder at a directory with one subfolder per class; the rest of the code stays the same. When you have only a few hundred images per class, stop training from scratch and fine-tune a pretrained model instead — freeze the convolutional layers, replace the final linear layer, and train that for a few epochs.

From there, add data augmentation, a learning-rate scheduler, and checkpointing that keeps the best validation score rather than the last epoch. Each of those is a twenty-minute change, and together they turn a toy into something you can actually rely on.

Photo: Daniil Komov / Pexels