Lesson 38: Transfer Learning

Training a good CNN from scratch (Lessons 34-35) took hundreds of labeled examples even for a toy task. Real target tasks are often data-starved: a handful of labeled medical scans, a new product category with 20 photos. Transfer learning sidesteps this by reusing a network already trained on a different, data-rich task, on the theory that early-layer features (edges, blobs, simple textures) are useful for almost any visual task, not just the one they were originally trained on. This lesson demonstrates it on real photos, using a CIFAR-10 subset (Krizhevsky, 2009).

In [1]:
import copy
import pickle
import tarfile
import urllib.request
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

Source task and target task

Set up two related but distinct tasks, both from CIFAR-10. The source task has plenty of data: distinguishing airplanes from automobiles, 300 training images each. The target task is the one we actually care about, and it's deliberately starved: distinguishing cats from trucks — two classes the source task never saw at all — with only 15 training images per class.

In [2]:
CIFAR_URL = 'https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'
CACHE_ROOT = Path.home() / '.cache' / 'cvintro'
CACHE_DIR = CACHE_ROOT / 'cifar-10-batches-py'

def ensure_cifar10():
    if CACHE_DIR.exists():
        return
    CACHE_ROOT.mkdir(parents=True, exist_ok=True)
    archive_path = CACHE_ROOT / 'cifar-10-python.tar.gz'
    if not archive_path.exists():
        print('Downloading CIFAR-10 (~163 MB, one-time, cached under ~/.cache/cvintro)...')
        urllib.request.urlretrieve(CIFAR_URL, archive_path)
    print('Extracting...')
    with tarfile.open(archive_path) as tar:
        tar.extractall(CACHE_ROOT)

def load_cifar_batch(path):
    with open(path, 'rb') as f:
        d = pickle.load(f, encoding='bytes')
    imgs = d[b'data'].reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1).astype(np.float32) / 255.0
    labels = np.array(d[b'labels'], dtype=np.int64)
    return imgs, labels

ensure_cifar10()

CIFAR_LABELS = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
CID = {name: CIFAR_LABELS.index(name) for name in CIFAR_LABELS}

train_imgs, train_labels = [], []
for i in range(1, 6):
    imgs, labels = load_cifar_batch(CACHE_DIR / f'data_batch_{i}')
    train_imgs.append(imgs); train_labels.append(labels)
train_imgs, train_labels = np.concatenate(train_imgs), np.concatenate(train_labels)
test_imgs, test_labels = load_cifar_batch(CACHE_DIR / 'test_batch')

def take(imgs, labels, name, n, rng_local):
    idx = np.where(labels == CID[name])[0]
    idx = rng_local.permutation(idx)[:n]
    return imgs[idx]

data_rng = np.random.default_rng(3)

# source task: airplane vs. automobile, plenty of data
src_a = take(train_imgs, train_labels, 'airplane', 300, data_rng)
src_b = take(train_imgs, train_labels, 'automobile', 300, data_rng)
X_src = np.concatenate([src_a, src_b])
y_src = np.concatenate([np.zeros(300), np.ones(300)]).astype(np.int64)
perm = data_rng.permutation(len(X_src)); X_src, y_src = X_src[perm], y_src[perm]

# target task: cat vs. truck, neither ever seen by the source task, deliberately starved
X_tgt_train = np.concatenate([take(train_imgs, train_labels, 'cat', 15, data_rng),
                               take(train_imgs, train_labels, 'truck', 15, data_rng)])
y_tgt_train = np.array([0] * 15 + [1] * 15, dtype=np.int64)
X_tgt_test = np.concatenate([take(test_imgs, test_labels, 'cat', 150, data_rng),
                              take(test_imgs, test_labels, 'truck', 150, data_rng)])
y_tgt_test = np.array([0] * 150 + [1] * 150, dtype=np.int64)

fig, axes = plt.subplots(2, 6, figsize=(11, 4))
for ax, im in zip(axes[0], list(src_a[:3]) + list(src_b[:3])):
    ax.imshow(im); ax.axis('off')
axes[0, 0].set_ylabel('source', rotation=0, labelpad=25)
for ax, im in zip(axes[1], list(X_tgt_train[:3]) + list(X_tgt_train[15:18])):
    ax.imshow(im); ax.axis('off')
axes[1, 0].set_ylabel('target', rotation=0, labelpad=25)
fig.suptitle('Source task (airplane vs. automobile, top) vs. target task (cat vs. truck, bottom)', y=1.02)
plt.show()
No description has been provided for this image

Three strategies

  1. From scratch — train a fresh CNN on only the 30 target images. This is the baseline: no transfer at all.
  2. Frozen backbone (feature extraction) — pretrain a CNN backbone on the source task, then freeze its weights entirely and train only a new linear classifier on top of the features it produces for target images.
  3. Fine-tuning — start from the same pretrained backbone, but keep updating it on the target data too, using a much smaller learning rate for the backbone than for the new classifier head (the backbone already encodes useful structure; large updates from just 30 examples would wreck it).
In [3]:
class Backbone(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 16, 5, padding=2), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 3, padding=1), nn.ReLU(),
            nn.AdaptiveMaxPool2d(1),
        )

    def forward(self, x):
        return self.conv(x).flatten(1)

class Classifier(nn.Module):
    def __init__(self, backbone, freeze_backbone):
        super().__init__()
        self.backbone = backbone
        self.freeze_backbone = freeze_backbone
        self.fc = nn.Linear(32, 2)

    def forward(self, x):
        feat = self.backbone(x)
        if self.freeze_backbone:
            feat = feat.detach()  # frozen via detach() here; the real ResNet below freezes via requires_grad=False instead -- both stop gradients, detach() is just cheaper when the backbone is never going to be unfrozen later
        return self.fc(feat)

def train_source(seed, epochs=300, lr=0.001):
    torch.manual_seed(seed)  # seed before constructing the model (Lesson 34/35)
    backbone = Backbone()
    fc = nn.Linear(32, 2)
    opt = torch.optim.Adam(list(backbone.parameters()) + list(fc.parameters()), lr=lr)
    Xt = torch.tensor(X_src).permute(0, 3, 1, 2); yt = torch.tensor(y_src)
    for _ in range(epochs):
        opt.zero_grad()
        loss = F.cross_entropy(fc(backbone(Xt)), yt)
        loss.backward()
        opt.step()
    return backbone

def train_target(backbone, freeze_backbone, seed, epochs=200, lr=0.01, backbone_lr=None):
    torch.manual_seed(seed)
    model = Classifier(backbone, freeze_backbone)
    if freeze_backbone:
        opt = torch.optim.Adam(model.fc.parameters(), lr=lr)
    elif backbone_lr is not None:
        opt = torch.optim.Adam([
            {'params': model.backbone.parameters(), 'lr': backbone_lr},
            {'params': model.fc.parameters(), 'lr': lr},
        ])
    else:
        opt = torch.optim.Adam(model.parameters(), lr=lr)
    Xtr = torch.tensor(X_tgt_train).permute(0, 3, 1, 2); ytr = torch.tensor(y_tgt_train)
    Xte = torch.tensor(X_tgt_test).permute(0, 3, 1, 2); yte = torch.tensor(y_tgt_test)
    for _ in range(epochs):
        opt.zero_grad()
        loss = F.cross_entropy(model(Xtr), ytr)
        loss.backward()
        opt.step()
    with torch.no_grad():
        return (model(Xte).argmax(1) == yte).float().mean().item()
In [4]:
pretrained = train_source(seed=0)  # pretrain once; the point of transfer learning is reusing this

scratch_accs, frozen_accs, finetune_accs = [], [], []
for seed in range(8):
    scratch_acc = train_target(Backbone(), freeze_backbone=False, seed=seed)
    frozen_acc = train_target(copy.deepcopy(pretrained), freeze_backbone=True, seed=seed)
    finetune_acc = train_target(copy.deepcopy(pretrained), freeze_backbone=False, seed=seed, backbone_lr=0.0003)
    scratch_accs.append(scratch_acc); frozen_accs.append(frozen_acc); finetune_accs.append(finetune_acc)

print(f'{"strategy":>20} {"mean test acc":>16} {"std":>8}')
print(f'{"from scratch":>20} {np.mean(scratch_accs):>15.1%} {np.std(scratch_accs):>8.1%}')
print(f'{"frozen backbone":>20} {np.mean(frozen_accs):>15.1%} {np.std(frozen_accs):>8.1%}')
print(f'{"fine-tuned":>20} {np.mean(finetune_accs):>15.1%} {np.std(finetune_accs):>8.1%}')
print(f'\n(averaged over {len(scratch_accs)} random seeds, each reusing the same 30 target training images)')
            strategy    mean test acc      std
        from scratch           63.6%     6.8%
     frozen backbone           69.5%     1.0%
          fine-tuned           76.7%     1.2%

(averaged over 8 random seeds, each reusing the same 30 target training images)

Both transfer strategies clearly beat training from scratch, and they're also dramatically more consistent (lower standard deviation). Even though the source task never saw a cat or a truck, it learned generic low-level features (edges, color blobs, textures) that transferred to the new task.

Note that fine-tuning used a learning rate for the backbone roughly 30x smaller than the classifier head's. With just 30 target examples, a larger backbone update would overfit those 30 images, discarding everything useful it had learned from the 600-image source task. This is the same catastrophic-forgetting failure mode as Lesson 35's overfitting curve, just applied to a network that started out already knowing something.

One more detail worth noting: pretrained is trained once (seed=0) and reused across all eight seeds in the loop above. Only the target-side training (classifier initialization, fine-tuning trajectory) varies with seed, so the standard deviations reported above measure variance from target-side training alone, not noise in the source pretraining.

In practice: pretrained backbones in one line

The Backbone above is trained from scratch on a toy task for illustrative purposes, but real transfer learning leverages existing networks available for reuse. torchvision ships architectures already pretrained on the full ImageNet dataset (Lesson 37) — 1.2 million images across 1,000 classes. Loading such a network, ready to use as a frozen feature extractor, is a single line of code:

In [5]:
import torchvision

resnet = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1)
n_total = sum(p.numel() for p in resnet.parameters())

resnet.fc = nn.Identity()  # strip the original 1000-class ImageNet head, keep the 512-d features underneath
resnet.eval()
for p in resnet.parameters():
    p.requires_grad = False  # frozen, feature-extraction style — the same strategy as "frozen backbone" above

print(f'total ResNet-18 parameters: {n_total:,}  (all pretrained on ImageNet, none of them touched here)')
total ResNet-18 parameters: 11,689,512  (all pretrained on ImageNet, none of them touched here)

ResNet-18 was trained on 224x224 images normalized by ImageNet's per-channel mean and standard deviation — not the plain [0, 1] pixel scaling used everywhere else in this lesson — so the target images need to be resized and renormalized before going through it. Extract a 512-dimensional feature vector for every target image once (the backbone is frozen, so this only needs to happen once, not on every training step), then train a small linear head on top — exactly the frozen-backbone strategy above, just with a very different backbone.

In [6]:
IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)

def extract_features(X):
    x = torch.tensor(X).permute(0, 3, 1, 2)
    x = F.interpolate(x, size=224, mode='bilinear', align_corners=False)
    x = (x - IMAGENET_MEAN) / IMAGENET_STD
    with torch.no_grad():
        return torch.cat([resnet(x[i:i + 32]) for i in range(0, len(x), 32)])

feat_train = extract_features(X_tgt_train)
feat_test = extract_features(X_tgt_test)
ytr_t = torch.tensor(y_tgt_train); yte_t = torch.tensor(y_tgt_test)

resnet_accs = []
for seed in range(8):
    torch.manual_seed(seed)
    head = nn.Linear(feat_train.shape[1], 2)
    opt = torch.optim.Adam(head.parameters(), lr=0.01)
    for _ in range(200):
        opt.zero_grad()
        loss = F.cross_entropy(head(feat_train), ytr_t)
        loss.backward()
        opt.step()
    with torch.no_grad():
        resnet_accs.append((head(feat_test).argmax(1) == yte_t).float().mean().item())

print(f'frozen real ImageNet ResNet-18: mean test acc = {np.mean(resnet_accs):.1%}  (+/- {np.std(resnet_accs):.1%}, 8 seeds)')
print(f'(toy frozen backbone above got {np.mean(frozen_accs):.1%})')
frozen real ImageNet ResNet-18: mean test acc = 98.1%  (+/- 0.2%, 8 seeds)
(toy frozen backbone above got 69.5%)

Frozen ImageNet features push accuracy well past either toy strategy above, using nothing but a new linear layer trained on the same 30 target images — no backbone training at all.

A caveat on this comparison: some of that gap is an unfair advantage, not just richer pretraining. ImageNet's 1,000 classes already include photos of cats and several vehicle types, unlike the toy source task, which only saw airplanes and automobiles — so this isn't a perfectly controlled comparison of "big general pretraining" versus "tiny 2-class pretraining" in isolation; part of ResNet-18's edge comes from having directly seen the target classes during pretraining.

Nevertheless, the qualitative point holds, and it's the actual reason ImageNet-pretrained backbones are the default starting point in practice: a network trained on 1.2 million images learns far richer, more broadly transferable features than one trained on a couple thousand images of two categories ever could — which is exactly what Lesson 37's architectures, and the vastly larger datasets they were trained on, are built to produce.

Exercises

  1. Try backbone_lr=0.01 (i.e. no learning-rate difference between backbone and head) in the fine-tuning call. Does fine-tuned accuracy get better or worse, and does that match the catastrophic-forgetting explanation above?
  2. Try freezing most of the backbone but fine-tuning only its last conv layer (hint: set requires_grad = False on the first Conv2d's parameters only). Where does that land relative to fully-frozen and fully-fine-tuned?
  3. Swap the target task's classes from cat/truck to dog/automobile — a target pair that overlaps the source task's domain (vehicles) even more closely than cat/truck did. Does frozen-backbone accuracy improve further, and does that match the intuition that transfer works best when source and target share more underlying visual structure?
  4. resnet.fc = nn.Identity() keeps every one of ResNet-18's convolutional layers frozen. Try unfreezing just its last residual stage (resnet.layer4) and fine-tuning it on the 30 target images at a small learning rate, the same way the toy backbone was fine-tuned above. Does it beat the fully-frozen ResNet-18 result, or is 30 images too few to safely update even a small fraction of 11 million parameters?
Take the Lesson 38 Quiz →