Lesson 36: Transfer Learning¶
Training a good CNN from scratch (Lessons 33-34) 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 — Lesson 33's Sobel-like first-layer filters) are useful for almost any visual task, not just the one they were originally trained on.
import copy
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. The source task has plenty of data: distinguishing plusses from circles, 300 training images. The target task is the one we actually care about, and it's deliberately starved: distinguishing squares from circles (a new class the source task never saw), with only 12 noisy training images.
def make_image(shape_type, cx, cy, size=16):
img = np.zeros((size, size), dtype=np.float32)
if shape_type == 'plus':
img[cy-1:cy+2, cx-3:cx+4] = 1.0
img[cy-3:cy+4, cx-1:cx+2] = 1.0
elif shape_type == 'circle':
yy, xx = np.mgrid[0:size, 0:size]
img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0
elif shape_type == 'square':
img[cy-3:cy+4, cx-3:cx+4] = 1.0
return img
def make_dataset(rng_local, n, shapes, position_range=(5, 11), noise=0.15):
imgs, labels = [], []
for _ in range(n):
shape_type = rng_local.choice(shapes)
cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)
img = make_image(shape_type, cx, cy)
img = np.clip(img + rng_local.normal(0, noise, img.shape), 0, 1).astype(np.float32)
imgs.append(img)
labels.append(shapes.index(shape_type))
return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.int64)
data_rng = np.random.default_rng(3)
X_src, y_src = make_dataset(data_rng, 300, ['plus', 'circle'], noise=0.15)
X_tgt_train, y_tgt_train = make_dataset(data_rng, 12, ['square', 'circle'], noise=0.4)
X_tgt_test, y_tgt_test = make_dataset(data_rng, 150, ['square', 'circle'], noise=0.4)
fig, axes = plt.subplots(2, 6, figsize=(11, 4))
for ax, im in zip(axes[0], X_src[:6]):
ax.imshow(im, cmap='gray'); ax.axis('off')
axes[0, 0].set_ylabel('source', rotation=0, labelpad=25)
for ax, im in zip(axes[1], X_tgt_train[:6]):
ax.imshow(im, cmap='gray'); ax.axis('off')
axes[1, 0].set_ylabel('target', rotation=0, labelpad=25)
fig.suptitle('Source task (plus vs. circle, top) vs. target task (square vs. circle, bottom)', y=1.02)
plt.show()
Three strategies¶
- From scratch — train a fresh CNN on only the 12 target images. This is the baseline: no transfer at all.
- 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.
- 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 12 examples would wreck it).
class Backbone(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(16, 32, 5, padding=2), 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() # no gradient flows into a frozen backbone
return self.fc(feat)
def train_source(seed, epochs=300, lr=0.01):
torch.manual_seed(seed) # seed before constructing the model (Lesson 33/34)
backbone = Backbone()
fc = nn.Linear(32, 2)
opt = torch.optim.Adam(list(backbone.parameters()) + list(fc.parameters()), lr=lr)
Xt = torch.tensor(X_src).unsqueeze(1); 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).unsqueeze(1); ytr = torch.tensor(y_tgt_train)
Xte = torch.tensor(X_tgt_test).unsqueeze(1); 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()
scratch_accs, frozen_accs, finetune_accs = [], [], []
for seed in range(8):
scratch_acc = train_target(Backbone(), freeze_backbone=False, seed=seed)
pretrained = train_source(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.0005)
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 12 target training images)')
Both transfer strategies beat training from scratch on average, and both are also noticeably more consistent (lower standard deviation) — with only 12 training images, a from-scratch network's success or failure depends heavily on which 12 images it happened to get, while a pretrained backbone starts from a much better place regardless. The source task never saw a square, yet the low-level features it learned (edges, curvature, blob-like regions) transferred anyway, because those features are generic to shape recognition, not specific to "plus vs. circle."
Note what fine-tuning needed to work at all: a backbone learning rate roughly 20x smaller than the classifier head's. With only 12 examples, an unrestrained backbone update would simply overfit those 12 images from scratch, discarding everything useful it learned from the 300-image source task — the same catastrophic-forgetting failure mode as Lesson 34's overfitting curve, just applied to a network that started out already knowing something.
Exercise¶
- 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? - Try freezing most of the backbone but fine-tuning only its last conv layer (hint: set
requires_grad = Falseon the firstConv2d's parameters only). Where does that land relative to fully-frozen and fully-fine-tuned? - The source and target tasks here share a class (circle). Design a source task that shares no classes with the target task at all, and predict whether transfer would still help. Test your prediction.