Hands-onIntermediate 40 min· Deep learning and modern AI · 17 of 20

Computer vision with CNNs and transfer learning

Learn how convolutional networks see images, then build an image classifier by reusing a pretrained ResNet with transfer learning. Free to run on a Colab or Kaggle GPU.

What you will build: an image classifier that tells apart 10 kinds of objects (aeroplane, car, bird, cat and so on) from the CIFAR-10 dataset, using a pretrained ResNet-18 and only a few minutes of training.

Prerequisites: *Deep learning with PyTorch: your first model*. Use a free Colab or Kaggle notebook with a GPU switched on (Colab: Runtime, Change runtime type; Kaggle: Settings, Accelerator). It also works on CPU but is slow.

First we look at how a CNN sees an image, then we skip the expensive part using transfer learning.

How a CNN sees an image

A colour image is a grid of numbers: height x width x 3 colour channels (red, green, blue). A fully connected network would need millions of weights for even a small photo and would ignore the fact that nearby pixels belong together.

A convolutional neural network (CNN) uses small filters (for example 3x3) that slide across the image. A filter is a tiny grid of weights that responds to one pattern: an edge, a corner, a patch of colour. Early layers find edges, middle layers combine them into textures and parts (an ear, a wheel), and later layers recognise whole objects.

The building blocks:

- Convolution layer (nn.Conv2d): applies many filters and produces feature maps. - ReLU: the usual activation. - Pooling layer (nn.MaxPool2d): shrinks the feature maps, keeping the strongest responses. - Fully connected layer at the end: turns the features into class scores.

Notebook cell

import torch
import torch.nn as nn

# PyTorch images are (batch, channels, height, width)
x = torch.randn(1, 3, 32, 32)            # one fake 32x32 colour image

conv = nn.Conv2d(in_channels=3, out_channels=8, kernel_size=3, padding=1)
pool = nn.MaxPool2d(2)

a = conv(x)
print("after conv:", a.shape)            # 8 feature maps, same size
b = pool(torch.relu(a))
print("after pool:", b.shape)            # height and width halved

Shapes are the best way to understand CNNs: printing them is also how you debug.

Output

after conv: torch.Size([1, 8, 32, 32])
after pool: torch.Size([1, 8, 16, 16])

A small CNN from scratch

small_cnn.py

class SmallCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),   # 32 -> 16
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),  # 16 -> 8
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, 64), nn.ReLU(),
            nn.Linear(64, num_classes),
        )

    def forward(self, x):
        return self.classifier(self.features(x))

print(SmallCNN()(torch.randn(4, 3, 32, 32)).shape)   # torch.Size([4, 10])

A model you write as a class has two methods: __init__ builds the layers and forward says how data flows through them. Trained from scratch on CIFAR-10 this kind of network learns something useful but needs many epochs.

Training your own CNN from scratch needs a lot of images and time. Most real projects do not do that. They use transfer learning.

Transfer learning

Researchers have already trained big CNNs on millions of images (the ImageNet dataset). Those networks learned general visual skills: edges, textures, shapes. You can borrow that knowledge.

The analogy: instead of teaching someone to see from birth, you hire someone who can already see and only teach them your 10 categories.

The recipe:

1. Load a pretrained model. 2. Freeze its layers so their weights do not change. 3. Replace the last layer with a new one that has your number of classes. 4. Train only the new layer (and later, optionally, fine-tune more layers at a low learning rate).

The plan

1Turn on a GPU in your notebook settings.
2Load CIFAR-10 with the pretrained model's own image transforms (Step 1).
3Load ResNet-18, freeze it and swap the last layer (Step 2).
4Train the new layer for a few epochs (Step 3).
5Evaluate and predict on single images (Step 4 and the next section).

Step 1: load and prepare the data

Pretrained models expect images prepared in a certain way (size 224x224 and specific colour normalisation). The weights object gives you the exact transforms, so you never guess them. To keep training quick we use 5,000 training images and 1,000 test images. The first run downloads CIFAR-10, which is about 170 MB.

Notebook cell 1

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, models

device = "cuda" if torch.cuda.is_available() else "cpu"
print("using", device)

weights = models.ResNet18_Weights.DEFAULT
preprocess = weights.transforms()          # resize, crop, convert to tensor, normalise

train_full = datasets.CIFAR10(root="data", train=True, download=True, transform=preprocess)
test_full = datasets.CIFAR10(root="data", train=False, download=True, transform=preprocess)

train_ds = Subset(train_full, range(5000))
test_ds = Subset(test_full, range(1000))

train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=64)

classes = train_full.classes
print(classes)

The class names are airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck. CIFAR-10 is already in random order, so the first 5,000 images contain all classes.

Step 2: load ResNet-18 and replace the last layer

Notebook cell 2

model = models.resnet18(weights=weights)

for param in model.parameters():
    param.requires_grad = False            # freeze everything

model.fc = nn.Linear(model.fc.in_features, 10)   # new head; its parameters are trainable
model = model.to(device)

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

Only the new final layer (512 x 10 weights plus 10 biases = 5,130 numbers) will learn. The other ~11 million stay fixed.

Output

trainable parameters: 5,130 of 11,181,642

Step 3: train the new head

Notebook cell 3

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.001)

for epoch in range(3):
    model.train()
    running = 0.0
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        optimizer.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()
        running += loss.item() * xb.size(0)
    print(f"epoch {epoch + 1}  train loss {running / len(train_ds):.4f}")

Multi-class problems use CrossEntropyLoss, which takes raw scores (logits) and integer class labels. No softmax in the model.

Step 4: evaluate and predict

Notebook cell 4

def evaluate(loader):
    model.eval()
    correct = 0
    with torch.no_grad():
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            preds = model(xb).argmax(dim=1)
            correct += (preds == yb).sum().item()
    return correct / len(loader.dataset)

print("test accuracy:", evaluate(test_loader))

# Predict one image with probabilities
img, label = test_ds[0]
model.eval()
with torch.no_grad():
    probs = torch.softmax(model(img.unsqueeze(0).to(device)), dim=1)[0]
top = probs.argmax().item()
print("actual:", classes[label], "| predicted:", classes[top], f"({probs[top].item():.2f})")

With a frozen ResNet-18 and 5,000 images, accuracy is usually well above what a small CNN gets in the same time (commonly somewhere in the 70s to 80s percent; results vary). Training from scratch for the same few minutes would do much worse.

argmax(dim=1) picks the class with the highest score. softmax turns the scores into probabilities that add up to 1, handy for showing confidence.

Predicting your own photo

Notebook cell 5

from PIL import Image

img = Image.open("my_photo.jpg").convert("RGB")   # upload a file to the notebook first
x = preprocess(img).unsqueeze(0).to(device)

model.eval()
with torch.no_grad():
    probs = torch.softmax(model(x), dim=1)[0]

for p, idx in zip(*probs.topk(3)):
    print(f"{classes[idx.item()]:12s} {p.item():.2f}")

Shows the top 3 guesses. Remember the model only knows these 10 classes, so a photo of a mango will still get one of them.

Fine-tuning and other tricks

Next steps

Technique	What it does	When to use
Feature extraction (what we did)	Freeze all, train new head	Small dataset, little time
Fine-tuning	Unfreeze the last block(s) and train with a very low learning rate such as 1e-4	You have more data and want extra accuracy
Data augmentation	Random flips, crops and colour changes on training images	Small datasets, to reduce overfitting
Bigger or newer backbones	ResNet-50, EfficientNet and others in torchvision	When accuracy matters more than speed
Your own folders	datasets.ImageFolder('data/train') with one subfolder per class	Your own image project

Check your data

A model is only as good as its data. CIFAR-10 pictures are tiny and centred. A model trained on them may fail on phone photos taken in poor light. Always test on images that look like your real use case, and watch for bias: if one group or setting is missing from training data, the model will perform worse for it.

Prompt · Plan an image classifier

I want to build an image classifier for [your problem, e.g. identifying 5 types of leaf diseases from phone photos]. I have about [number] images and a free Colab GPU. Suggest: how to organise the folders for torchvision ImageFolder, which pretrained model to start with, what data augmentation to use, how to split train/validation/test, and 3 ways my model could fail in real use.

What to expect: Check that it suggests a validation split, a baseline, and real-world failure cases like lighting and background. Verify any torchvision API names in the official docs.

Try it: Train your own classifier

Download a small image dataset from Kaggle (for example a flowers or cats-versus-dogs dataset), arrange it as one folder per class, load it with datasets.ImageFolder(root, transform=preprocess), split it 80/20 with torch.utils.data.random_split, and reuse this lesson's code. Then unfreeze model.layer4 and train two more epochs with lr=1e-4 for the unfrozen layers. Did accuracy improve?

Key takeaways

CNNs use sliding filters to detect patterns, building from edges to whole objects.
Printing tensor shapes (batch, channels, height, width) is the key to understanding and debugging a CNN.
Transfer learning reuses a pretrained model: freeze it, replace the last layer, train only the new head.
Use weights.transforms() so your images are prepared exactly as the pretrained model expects.
Fine-tune deeper layers only with a low learning rate and enough data.
Test on images that resemble real use, not just the benchmark dataset.

Check yourself

Finishing the quiz marks the lesson as done.

1. Why are CNNs better than fully connected networks for images?

2. In transfer learning, what does 'freezing' layers mean?

3. You have only 800 labelled photos for a new 4-class problem. What is the best starting approach?

4. Which loss function fits a 10-class image problem with integer labels and raw model outputs?

5. An input batch has shape (64, 3, 224, 224). What does each number mean?

Saved on this device. Create a free account to keep it on every device.

Deep learning with PyTorch: your first model NLP and transformers with Hugging Face

Practise this with Sam

A real mock interview, out loud, with an honest scorecard. Free to start.

Start free