Lesson 34: Training a CNN¶
Lesson 33's CNN was trained and evaluated on data drawn from the same, fairly generous distribution. Real training has a much sharper failure mode lurking: with too little data and too much model capacity, a network can perfectly memorize its training set while learning nothing that generalizes. This lesson makes that failure concrete, then fixes it two different ways: data augmentation (Lesson 8's transforms, repurposed) and regularization (weight decay).
import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
A deliberately hard, small dataset¶
The same plus-vs-circle task as Lesson 33, but now with only 12 training images and pixel noise added to every image, while the validation set stays large (150 images) so its accuracy is a reliable estimate.
def make_image(shape_type, cx, cy, size=16, rng=None):
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
else:
yy, xx = np.mgrid[0:size, 0:size]
img[((xx - cx) ** 2 + (yy - cy) ** 2) <= 9] = 1.0
if rng is not None:
img = np.clip(img + rng.normal(0, 0.4, img.shape), 0, 1).astype(np.float32)
return img
def make_dataset(rng_local, n, position_range):
imgs, labels = [], []
for _ in range(n):
shape_type = rng_local.choice(['plus', 'circle'])
cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)
imgs.append(make_image(shape_type, cx, cy, rng=rng_local))
labels.append(0.0 if shape_type == 'plus' else 1.0)
return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.float32)
data_rng = np.random.default_rng(2)
X_train, y_train = make_dataset(data_rng, 12, (3, 13))
X_val, y_val = make_dataset(data_rng, 150, (3, 13))
fig, axes = plt.subplots(1, 6, figsize=(11, 2))
for ax, im in zip(axes, X_train[:6]):
ax.imshow(im, cmap='gray')
ax.axis('off')
fig.suptitle('The entire training set is only 12 noisy images like these', y=1.05)
plt.show()
Watching it overfit¶
Train a reasonably large CNN (Lesson 33's architecture) on just these 12 images, and track both training loss and validation loss (on the held-out 150 images) at every epoch.
class CNN(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),
)
self.fc = nn.Linear(32, 1)
def forward(self, x):
return self.fc(self.conv(x).flatten(1)).squeeze(-1)
def train_tracked(model_cls, Xtr, ytr, Xval, yval, epochs=400, lr=0.01, weight_decay=0.0, seed=0):
torch.manual_seed(seed) # seed BEFORE constructing the model, so init is actually reproducible
model = model_cls()
opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
train_losses, val_losses = [], []
for _ in range(epochs):
opt.zero_grad()
loss = F.binary_cross_entropy_with_logits(model(Xtr), ytr)
loss.backward()
opt.step()
with torch.no_grad():
train_losses.append(loss.item())
val_losses.append(F.binary_cross_entropy_with_logits(model(Xval), yval).item())
with torch.no_grad():
train_acc = ((model(Xtr) > 0).float() == ytr).float().mean().item()
val_acc = ((model(Xval) > 0).float() == yval).float().mean().item()
return train_losses, val_losses, train_acc, val_acc
Xtr_t = torch.tensor(X_train).unsqueeze(1); ytr_t = torch.tensor(y_train)
Xval_t = torch.tensor(X_val).unsqueeze(1); yval_t = torch.tensor(y_val)
train_losses, val_losses, train_acc, val_acc = train_tracked(CNN, Xtr_t, ytr_t, Xval_t, yval_t)
print(f'final train accuracy: {train_acc:.1%}')
print(f'final val accuracy: {val_acc:.1%}')
print(f'val loss minimum was {min(val_losses):.3f} at epoch {np.argmin(val_losses)} '
f'(out of {len(val_losses)}); it ended at {val_losses[-1]:.3f}')
plt.plot(train_losses, label='train loss')
plt.plot(val_losses, label='val loss')
plt.axvline(np.argmin(val_losses), color='gray', linestyle='--', linewidth=1, label='best val loss')
plt.xlabel('epoch'); plt.ylabel('loss'); plt.legend(fontsize=8)
plt.title('The classic overfitting curve')
plt.show()
Training loss marches steadily to zero — the network perfectly memorizes all 12 images, noise included. Validation loss, meanwhile, bottoms out part-way through training and then climbs back up: past that point, every further epoch makes the model more confidently wrong about data it hasn't seen. Final validation accuracy lands well short of the training set's perfect score, despite the training set being fit exactly.
Fix 1: data augmentation¶
If there isn't enough real data, manufacture more from what's there. Apply random transformations from Lesson 8 (flips, small rotations) to each training image — the label doesn't change, but the pixels do, so the network sees a much wider variety of "what a plus/circle can look like" instead of memorizing 12 exact images.
def augment(img, rng_local):
if rng_local.random() < 0.5:
img = np.fliplr(img).copy()
if rng_local.random() < 0.5:
img = np.flipud(img).copy()
angle = rng_local.uniform(-10, 10)
M = cv2.getRotationMatrix2D((8, 8), angle, 1.0)
img = cv2.warpAffine(img, M, (16, 16))
return img.astype(np.float32)
aug_rng = np.random.default_rng(5)
X_aug, y_aug = [], []
for _ in range(20): # 20 augmented copies of each of the 12 original images
for img, label in zip(X_train, y_train):
X_aug.append(augment(img, aug_rng))
y_aug.append(label)
X_aug, y_aug = np.array(X_aug, dtype=np.float32), np.array(y_aug, dtype=np.float32)
fig, axes = plt.subplots(1, 6, figsize=(11, 2))
for ax, im in zip(axes, X_aug[:6]):
ax.imshow(im, cmap='gray')
ax.axis('off')
fig.suptitle(f'6 of {len(X_aug)} augmented copies, all still "the same 12 base images"', y=1.05)
plt.show()
Xaug_t = torch.tensor(X_aug).unsqueeze(1); yaug_t = torch.tensor(y_aug)
_, _, aug_train_acc, aug_val_acc = train_tracked(CNN, Xaug_t, yaug_t, Xval_t, yval_t, epochs=150)
print(f'with augmentation: train acc = {aug_train_acc:.1%}, val acc = {aug_val_acc:.1%} (was {val_acc:.1%})')
Fix 2: weight decay¶
Weight decay adds a penalty proportional to the squared weight magnitudes directly into the loss (equivalently, it shrinks every weight slightly toward zero on every update). Large, highly-tuned weights are exactly what a network needs to memorize 12 specific noisy images; penalizing weight magnitude makes that memorization more costly relative to finding a simpler, smoother function — without adding a single extra training example.
_, _, wd_train_acc, wd_val_acc = train_tracked(CNN, Xtr_t, ytr_t, Xval_t, yval_t, weight_decay=0.05)
print(f'with weight_decay=0.05: train acc = {wd_train_acc:.1%}, val acc = {wd_val_acc:.1%} (was {val_acc:.1%})')
print()
print(f'{"approach":>20} {"train acc":>12} {"val acc":>12}')
print(f'{"no fix":>20} {train_acc:>12.1%} {val_acc:>12.1%}')
print(f'{"augmentation":>20} {aug_train_acc:>12.1%} {aug_val_acc:>12.1%}')
print(f'{"weight decay":>20} {wd_train_acc:>12.1%} {wd_val_acc:>12.1%}')
Both fixes recover a large chunk of the lost validation accuracy, from two different angles: augmentation attacks the problem by giving the model more (synthetic) data to be right about; weight decay attacks it by making the model less willing to contort itself around a small dataset in the first place. In practice, both are normally used together, along with other regularizers like dropout, covered next, and batch normalization (Lesson 35), which incidentally also acts as a mild regularizer.
Fix 3: dropout¶
Dropout (Srivastava et al., 2014★) randomly zeroes out a fraction p of a layer's activations on every training step, forcing the surviving units to not rely on any one specific other unit always being present. To keep the layer's output at the same overall scale whether or not dropout is active, the surviving activations are rescaled by 1 / (1 - p) — this is "inverted dropout," what every framework's Dropout layer actually implements. At evaluation time, dropout does nothing at all: the full, unmodified layer runs, which is why model.eval() (used throughout this lesson already, for weight decay and augmentation too) matters — forgetting it would leave dropout randomly firing at test time.
def manual_dropout(x, p, rng_gen):
keep_prob = 1 - p
mask = (torch.rand(x.shape, generator=rng_gen) < keep_prob).float()
return x * mask / keep_prob # rescale so E[output] == input
x_demo = torch.randn(2000, 10)
torch.manual_seed(0)
out_torch = F.dropout(x_demo, p=0.3, training=True)
out_manual = manual_dropout(x_demo, 0.3, torch.Generator().manual_seed(0))
print(f'fraction zeroed, torch: {(out_torch == 0).float().mean().item():.3f} (target p = 0.3)')
print(f'fraction zeroed, manual: {(out_manual == 0).float().mean().item():.3f}')
print(f'mean before dropout: {x_demo.mean().item():.4f}')
print(f'mean after dropout (torch): {out_torch.mean().item():.4f} (rescaling keeps this close to the input mean)')
print(f'mean after dropout (manual): {out_manual.mean().item():.4f}')
eval_mode_out = F.dropout(x_demo, p=0.3, training=False)
print(f'eval-mode dropout is a no-op: {torch.equal(eval_mode_out, x_demo)}')
Now apply it to this lesson's overfitting problem. Dropout needs some redundancy to work with — zeroing half of a 32-unit feature vector going straight into a 1-unit output leaves little room to help, so add a wider hidden layer (32 → 64 → 1) and place dropout on the 64-unit layer. With only 12 training images, results are noisy from one random seed to the next, so compare mean validation accuracy over several seeds rather than trusting a single run.
class CNNDropout(nn.Module):
def __init__(self, dropout_p=0.0):
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),
)
self.fc1 = nn.Linear(32, 64)
self.dropout = nn.Dropout(dropout_p)
self.fc2 = nn.Linear(64, 1)
def forward(self, x):
feat = torch.relu(self.fc1(self.conv(x).flatten(1)))
return self.fc2(self.dropout(feat)).squeeze(-1)
def val_acc_for(dropout_p, seed, epochs=150):
torch.manual_seed(seed)
model = CNNDropout(dropout_p)
opt = torch.optim.Adam(model.parameters(), lr=0.01)
for _ in range(epochs):
model.train()
opt.zero_grad()
loss = F.binary_cross_entropy_with_logits(model(Xtr_t), ytr_t)
loss.backward()
opt.step()
model.eval()
with torch.no_grad():
return ((model(Xval_t) > 0).float() == yval_t).float().mean().item()
no_drop_accs = [val_acc_for(0.0, seed) for seed in range(8)]
drop_accs = [val_acc_for(0.5, seed) for seed in range(8)]
print(f'no dropout: mean val acc = {np.mean(no_drop_accs):.1%} (+/- {np.std(no_drop_accs):.1%}, 8 seeds)')
print(f'dropout(0.5): mean val acc = {np.mean(drop_accs):.1%} (+/- {np.std(drop_accs):.1%}, 8 seeds)')
Dropout gives a modest average improvement here, though the gap is small relative to the run-to-run noise from having only 12 training images — a much less dramatic effect than augmentation or weight decay showed above. That's realistic: dropout's benefit is well established at the scale of real datasets and real networks (hundreds of redundant units, thousands of examples), but in a toy setting this tiny, there just isn't much redundancy for it to exploit yet. In practice it is almost always combined with the other regularizers on this page, not used alone.
Exercise¶
- Try
weight_decayvalues of0.001,0.05, and1.0. Is there a point where it starts to hurt training accuracy along with (eventually) validation accuracy? What does an excessively large weight decay do to the model's capacity to fit anything at all? - Sweep
dropout_pover[0.0, 0.2, 0.4, 0.6, 0.8]inval_acc_for, averaging over the same 8 seeds at each value. Is there a value that's clearly best, or does the mean stay within one standard deviation across most of the range given how little data there is? - Increase the augmentation multiplier from 20 to 100 copies per base image. Does validation accuracy keep improving, or does it plateau — and if it plateaus, what does that suggest about the fundamental limit of augmenting a dataset that only contains 12 underlying examples to begin with?