Data Augmentation
More data for free.
More data is the best cure for overfitting. But collecting and labelling images is expensive. Here's a neat trick. A cat flipped from left to right is still a cat. Slightly cropped, or a little brighter, it's still a cat. So from one photo, we can make endless new training examples, for free.
- More data costs money
- A flipped cat is still a cat
- Endless examples for free
One photo, many views
Torchvision has a toolbox of these transforms. We compose three. A random flip, a random crop, and random colour changes. Run it. Here's one original photo, and five views of it. Each time an image is loaded for training, the model gets a slightly different version. So it can't memorise. It has to learn what really makes this a truck. Note that we only do this for training, never for testing.
import torch
from torchvision import datasets, transforms
from torchvision.transforms import v2
import matplotlib.pyplot as plt
image = datasets.CIFAR10("data", train=False, download=True)[3][0]
aug = v2.Compose([v2.RandomHorizontalFlip(),
v2.RandomResizedCrop(32, scale=(0.6, 1.0)),
v2.ColorJitter(0.4, 0.4, 0.4)])
fig, ax = plt.subplots(1, 6, figsize=(10, 2))
ax[0].imshow(image); ax[0].set_title("Original")
torch.manual_seed(1)
for a in ax[1:]:
a.imshow(aug(image))
for a in ax: a.axis("off")
plt.tight_layout()
plt.savefig("augment.png", dpi=120)
print("Five new training images from one photo")Five new training images from one photo
Choosing augmentations
Think about what stays true. Flipping is fine for animals, but not for digits, where a flipped two isn't a two. Rotation is fine for satellite images, but a bit silly for street signs. Colour jitter is fine, unless colour is the answer. Strong modern recipes like RandAugment, MixUp, and CutMix go much further, and blend images together. Always ask, does the label stay true?
A real image classifier
Let's bring this all together, on real colour photos of planes, cars, birds, and more. The model uses three blocks with batch norm, global pooling, and dropout. The data uses crops and flips. We train on twelve thousand images, for speed. Run it. Accuracy climbs steadily, on a hard ten class problem, with just a few epochs and a small network. Aria can now recognise real objects in real photos. That's a big milestone.
import torch
from torch import nn
from torchvision import datasets
from torchvision.transforms import v2
root = "data"
norm = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True),
v2.Normalize([0.49, 0.48, 0.45], [0.25, 0.24, 0.26])])
aug = v2.Compose([v2.RandomCrop(32, padding=4), v2.RandomHorizontalFlip(), norm])
train = datasets.CIFAR10(root, train=True, download=True, transform=aug)
test = datasets.CIFAR10(root, train=False, download=True, transform=norm)
train = torch.utils.data.Subset(train, range(12000))
test = torch.utils.data.Subset(test, range(2000))
def block(i, o):
return nn.Sequential(nn.Conv2d(i, o, 3, padding=1), nn.BatchNorm2d(o),
nn.ReLU(), nn.MaxPool2d(2))
model = nn.Sequential(block(3, 32), block(32, 64), block(64, 128),
nn.AdaptiveAvgPool2d(1), nn.Flatten(),
nn.Dropout(0.2), nn.Linear(128, 10))
opt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=0.01)
dl = torch.utils.data.DataLoader(train, batch_size=128, shuffle=True)
tl = torch.utils.data.DataLoader(test, batch_size=500)
for epoch in range(1, 7):
model.train()
for x, y in dl:
loss = nn.functional.cross_entropy(model(x), y)
opt.zero_grad(); loss.backward(); opt.step()
if epoch % 2 == 0:
model.eval()
hits = 0
with torch.no_grad():
for x, y in tl:
hits += (model(x).argmax(1) == y).sum().item()
print(f"epoch {epoch} accuracy {hits / len(test):.3f}")epoch 2 accuracy 0.443 epoch 4 accuracy 0.533 epoch 6 accuracy 0.540
Summary
Augmentation gives you free data, and it's one of the strongest tools there is. Aria can now see. But there's a problem with sight alone. The world doesn't arrive as single pictures. Language, speech, and video all come in sequences. Next, teaching networks to remember.
Key takeaways- Augment training data only
- Ask: is the label still true?
- Aria can now see