Transfer Learning in Practice

The last lesson ended with a promise: every architecture in torchvision ships with weights already trained on ImageNet. In this lesson you'll cash that in. Instead of training a CNN from scratch, you'll take a pretrained ResNet-18, swap its final layer for one that fits your classes, and get a strong classifier from a few hundred images in minutes. All the PyTorch code here belongs in Colab with a GPU runtime — transfer learning is cheap, but not browser-cheap.

Why features transfer

Recall what a CNN learns layer by layer: the earliest convolutions detect edges, color blobs, and simple textures; middle layers combine those into motifs like fur, mesh, or eyes; only the last layers become specific to the 1,000 ImageNet categories. Here's the key observation — the early layers are generic. An edge detector trained on dogs and teapots works just as well on X-rays and satellite photos. Those generic features took 1.28 million labeled images and serious GPU time to learn. Your dataset of 2,000 photos can't reproduce them — but it doesn't have to, because you can download them.

Transfer learning reuses a pretrained network as a feature extractor and only re-learns the part that's actually specific to your problem: the classifier head.

Two strategies

There are two standard moves, usually applied in sequence:

• Feature extraction (adaptation). Load the pretrained model, freeze every parameter in the backbone, replace the final classification layer with a fresh one sized for your classes, and train only that new head. Fast, data-efficient, and hard to mess up — the pretrained weights can't be damaged because they never change.
• Fine-tuning. After the head has converged, unfreeze some or all of the backbone and keep training with a much lower learning rate — around 10× to 100× lower (say 1e-5 instead of 1e-3). The small steps gently adapt the pretrained features to your domain without destroying them. If it helps, repeat with an even lower rate.

Which strategy, and how much to unfreeze, depends on two questions: how much data do you have, and how similar is your domain to ImageNet?

More data or a more different domain both push you toward unfreezing more.

The pipeline: data

torchvision's ImageFolder turns a directory tree into a dataset — one subfolder per class, folder names become labels:

data/
  train/
    cats/  cat001.jpg  cat002.jpg  ...
    dogs/  dog001.jpg  ...
  test/
    cats/  ...
    dogs/  ...

One rule is non-negotiable: your inputs must look like what the network saw during pretraining. That means 224×224 crops and normalization with the ImageNet channel statistics — mean [0.485, 0.456, 0.406] and std [0.229, 0.224, 0.225]:

import torch
from torch import nn, optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
train_transform = transforms.Compose([
    transforms.RandomRotation(10),
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
 
test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
 
train_set = datasets.ImageFolder("data/train", transform=train_transform)
test_set = datasets.ImageFolder("data/test", transform=test_transform)
trainloader = DataLoader(train_set, batch_size=64, shuffle=True, num_workers=2)
testloader = DataLoader(test_set, batch_size=64)
 
num_classes = len(train_set.classes)
print(train_set.classes)

Augmentation (rotation, random crop, flip) goes on the training transform only — evaluation uses a deterministic resize and center crop.

What does that Normalize actually do? A quick browser check:

Each channel is shifted and scaled exactly the way ImageNet images were during pretraining. Skip this (or use different stats) and the frozen backbone receives inputs from a distribution it has never seen — accuracy quietly tanks and nothing errors out.

The pipeline: model

Load ResNet-18 with pretrained weights, freeze everything, then replace the head. In ResNet the head is a single layer called fc:

from torchvision.models import resnet18, ResNet18_Weights
 
model = resnet18(weights=ResNet18_Weights.DEFAULT)
 
for param in model.parameters():
    param.requires_grad = False          # freeze the backbone
 
model.fc = nn.Linear(model.fc.in_features, num_classes)   # new head, 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: {trainable:,} / {total:,}")

New layers are created with requires_grad=True by default, so only the fresh fc will learn. Out of 11 million parameters, you're training a few thousand.

Phase 1: train the head

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.fc.parameters(), lr=1e-3)
 
def run_epoch(loader, train=True):
    model.train() if train else model.eval()
    total_loss = correct = n = 0
    ctx = torch.enable_grad() if train else torch.inference_mode()
    with ctx:
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            output = model(images)
            loss = criterion(output, labels)
            if train:
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()
            total_loss += loss.item() * len(images)
            correct += (output.argmax(1) == labels).sum().item()
            n += len(images)
    return total_loss / n, correct / n
 
for epoch in range(5):
    train_loss, train_acc = run_epoch(trainloader, train=True)
    test_loss, test_acc = run_epoch(testloader, train=False)
    print(f"epoch {epoch}: train_acc={train_acc:.3f}  test_acc={test_acc:.3f}")

Because only the head learns, this converges in a handful of epochs.

Phase 2: fine-tune

Now unfreeze and continue with a much smaller learning rate. Give it more patience — improvements come slowly and gently:

for param in model.parameters():
    param.requires_grad = True           # unfreeze everything
 
optimizer = optim.AdamW(model.parameters(), lr=1e-5)   # ~100x lower
 
best_acc, patience, bad_epochs = 0.0, 3, 0
for epoch in range(20):
    run_epoch(trainloader, train=True)
    _, test_acc = run_epoch(testloader, train=False)
    if test_acc > best_acc:
        best_acc, bad_epochs = test_acc, 0
        torch.save(model.state_dict(), "resnet_best.pth")
    else:
        bad_epochs += 1
        if bad_epochs >= patience:
            break
    print(f"epoch {epoch}: test_acc={test_acc:.3f}  (best {best_acc:.3f})")

Why the tiny learning rate? The pretrained weights are already excellent — big steps would scramble them, and you'd be back to (badly) training from scratch. Fine-tuning is a polish, not a rebuild.

The benchmark mindset

Never trust a technique without a baseline. On a typical small custom dataset (a few thousand images, a handful of classes), the comparison looks like this:

The exact numbers depend on your data; the pattern is remarkably stable. Feature extraction alone usually captures most of the gain; fine-tuning adds a few extra points on top.

Practical tips

• Match the preprocessing. 224×224 inputs, ImageNet mean/std. Each weights object documents its own recipe: ResNet18_Weights.DEFAULT.transforms() returns the exact transform used at pretraining time.
• Head first, backbone second. Training the head with a frozen backbone first prevents large random-head gradients from wrecking pretrained weights.
• Unfreeze in proportion to your data. More images, or a domain further from ImageNet, justify unfreezing more layers. With tiny datasets, keep the backbone frozen.
• Lower the learning rate when you unfreeze. Rule of thumb: divide by 10 to 100. If fine-tuning makes things worse, your rate is too high.

Next module: images are done — we turn to data with order. Recurrent neural networks give a model memory, one time step at a time.