Lesson 35: Training a CNN¶
Lesson 34'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. Lesson 32 showed that more capacity means a model can represent more; this lesson shows the flip side — that same surplus capacity, with too little data to constrain it, is exactly what lets a network memorize instead of generalize. This lesson makes that failure concrete, then fixes it three different ways: data augmentation (Lesson 8's transforms, repurposed), regularization (weight decay), and dropout.
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 34, 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 34'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 37), 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)')
In a toy setting this tiny (just 12 training images), there isn't much redundancy for dropout to exploit yet. Although only a modest improvement is shown here, dropout's benefit is well established at the scale of real datasets and real networks (hundreds of redundant units, thousands of examples). In practice dropout is almost always combined with the other regularizers on this page, not used alone.
Putting it together: real photos, real overfitting¶
Every fix so far has been tested one at a time, on a synthetic 12-image toy problem chosen specifically to make the failure mode unmistakable. This closing section repeats the experiment on real CIFAR-10 (Krizhevsky, 2009) photographs: a small, deliberately data-scarce training set of cats and trucks, with augmentation, weight decay, and dropout combined together. Real photographs contain lighting, clutter, and genuine visual ambiguity between classes.
import pickle
import tarfile
import urllib.request
from pathlib import Path
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
CIFAR_LABELS = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
ensure_cifar10()
cifar_train_imgs, cifar_train_labels = [], []
for i in range(1, 6):
imgs, labels = load_cifar_batch(CACHE_DIR / f'data_batch_{i}')
cifar_train_imgs.append(imgs)
cifar_train_labels.append(labels)
cifar_train_imgs = np.concatenate(cifar_train_imgs)
cifar_train_labels = np.concatenate(cifar_train_labels)
cifar_test_imgs, cifar_test_labels = load_cifar_batch(CACHE_DIR / 'test_batch')
def take_class(imgs, labels, name, n, rng_local):
idx = np.where(labels == CIFAR_LABELS.index(name))[0]
idx = rng_local.permutation(idx)[:n]
return imgs[idx].copy()
cifar_rng = np.random.default_rng(4)
N_CIFAR_TRAIN = 40 # per class -- deliberately small, real training data is expensive
X_cifar_train = np.concatenate([
take_class(cifar_train_imgs, cifar_train_labels, 'cat', N_CIFAR_TRAIN, cifar_rng),
take_class(cifar_train_imgs, cifar_train_labels, 'truck', N_CIFAR_TRAIN, cifar_rng),
])
y_cifar_train = np.array([0.0] * N_CIFAR_TRAIN + [1.0] * N_CIFAR_TRAIN, dtype=np.float32)
X_cifar_val = np.concatenate([
take_class(cifar_test_imgs, cifar_test_labels, 'cat', 150, cifar_rng),
take_class(cifar_test_imgs, cifar_test_labels, 'truck', 150, cifar_rng),
])
y_cifar_val = np.array([0.0] * 150 + [1.0] * 150, dtype=np.float32)
fig, axes = plt.subplots(1, 6, figsize=(11, 2))
for ax, im in zip(axes, X_cifar_train[:6]):
ax.imshow(im)
ax.axis('off')
fig.suptitle(f'6 of {len(X_cifar_train)} real training photos (cat vs. truck)', y=1.05)
plt.show()
Same recipe as before, adapted to 3-channel input: a plain CNN with no regularization, versus the same architecture with augmentation, weight decay, and dropout all applied together in a single run. Each configuration is averaged over 3 seeds, since 80 training images is still small enough for individual runs to vary.
class CNNRGB(nn.Module):
def __init__(self, dropout_p=0.0):
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.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 augment_photo(img, rng_local):
# unlike `augment` above, no vertical flip: upside-down cats and trucks aren't realistic
if rng_local.random() < 0.5:
img = np.fliplr(img).copy()
angle = rng_local.uniform(-10, 10)
M = cv2.getRotationMatrix2D((16, 16), angle, 1.0)
img = cv2.warpAffine(img, M, (32, 32))
return img.astype(np.float32)
Xcv_t = torch.tensor(X_cifar_val).permute(0, 3, 1, 2); ycv_t = torch.tensor(y_cifar_val)
def train_cifar(use_aug, weight_decay, dropout_p, seed, epochs=300, lr=0.001):
torch.manual_seed(seed)
if use_aug:
local_aug_rng = np.random.default_rng(seed + 100)
Xa, ya = [], []
for _ in range(10): # 10 augmented copies of each of the 80 base images
for img, label in zip(X_cifar_train, y_cifar_train):
Xa.append(augment_photo(img, local_aug_rng))
ya.append(label)
Xtr_np, ytr_np = np.array(Xa, dtype=np.float32), np.array(ya, dtype=np.float32)
else:
Xtr_np, ytr_np = X_cifar_train, y_cifar_train
Xct_t = torch.tensor(Xtr_np).permute(0, 3, 1, 2); yct_t = torch.tensor(ytr_np)
model = CNNRGB(dropout_p=dropout_p)
opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
for _ in range(epochs):
model.train()
opt.zero_grad()
loss = F.binary_cross_entropy_with_logits(model(Xct_t), yct_t)
loss.backward()
opt.step()
model.eval()
with torch.no_grad():
train_acc = ((model(Xct_t) > 0).float() == yct_t).float().mean().item()
val_acc = ((model(Xcv_t) > 0).float() == ycv_t).float().mean().item()
return train_acc, val_acc
cifar_configs = [
('no fix', dict(use_aug=False, weight_decay=0.0, dropout_p=0.0)),
('augmentation', dict(use_aug=True, weight_decay=0.0, dropout_p=0.0)),
('weight decay', dict(use_aug=False, weight_decay=0.01, dropout_p=0.0)),
('dropout', dict(use_aug=False, weight_decay=0.0, dropout_p=0.3)),
('all combined', dict(use_aug=True, weight_decay=0.01, dropout_p=0.3)),
]
print(f'{"approach":>15} {"train acc":>12} {"val acc":>12}')
cifar_results = {}
for name, cfg in cifar_configs:
tr_accs, val_accs = [], []
for seed in range(3):
tr, va = train_cifar(seed=seed, **cfg)
tr_accs.append(tr)
val_accs.append(va)
cifar_results[name] = (np.mean(tr_accs), np.mean(val_accs))
print(f'{name:>15} {np.mean(tr_accs):>12.1%} {np.mean(val_accs):>12.1%}')
No single fix is a clear winner here. Weight decay alone actually lands below the unregularized baseline — with weight_decay=0.01, training accuracy no longer even reaches 100%, meaning the penalty is strong enough to fight the objective without finding a better solution in its place. That's a useful, humbling result: the weight_decay=0.05 value used earlier in this lesson was tuned for a much smaller network on 12 tiny synthetic images, and a regularization strength that works well in one setting doesn't automatically transfer to a different architecture or real photographic data.
Augmentation and dropout each recover a modest amount of validation accuracy individually.
"All combined" beats every single fix and the unregularized baseline, echoing the same lesson from the synthetic dataset: these three regularizers compound rather than substitute for each other. The gain over the unregularized baseline is real but modest (a few percentage points here), which is itself the honest takeaway: regularization narrows the gap between training and validation performance, it doesn't manufacture data or model capacity that isn't there.
Closing that gap further, on real images, takes what later lessons introduce — more training data, better architectures (Lesson 37), and eventually transfer learning (Lesson 38): starting from a network already trained on millions of images, instead of memorizing from scratch on a handful of them.
Exercises¶
- 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?