Transfer Learning

Training a great image model from nothing takes millions of images, and weeks of G P U time. So here's the best kept secret in practical deep learning. You almost never do that. Someone else has already done the hard work, and they've given it away for free. Most real projects start with a model that's already trained.

• Millions of images, weeks of GPU
• You almost never start from zero
• The best kept secret

Why it works

Remember how a network builds understanding in layers. Early layers learn edges and textures, and these are useful for almost any image. Only the last layers are specific to the original task. So we take a network trained on a million photos, cut off its final layer, and attach a new one for our task. The early layers bring years of visual experience for free. We only need to teach the last step.

Loading a trained giant

With a single line, we download a ResNet trained on over a million images. Run it. Over eleven million weights, already learned, for a thousand kinds of object. And there at the end is the layer we'll replace, called f c. It maps five hundred and twelve features to a thousand classes. Getting all of this for free is remarkable, and it's the reason that small teams can build amazing things.

import torch
from torch import nn
from torchvision import datasets, models

weights = models.ResNet18_Weights.DEFAULT
model = models.resnet18(weights=weights)

print("Weights already learned:", sum(p.numel() for p in model.parameters()))
print("Trained on:", len(weights.meta["categories"]), "object types")
print("Last layer:", model.fc)

Weights already learned: 11689512
Trained on: 1000 object types
Last layer: Linear(in_features=512, out_features=1000, bias=True)

Freezing and replacing

Now let's use it on our own task, with just twelve hundred images. We freeze every weight, by turning off their gradients. We replace the final layer with a fresh one, for ten classes. The optimiser only trains that new layer. Run it. Remember that our own network, trained on twelve thousand images, took several epochs to do far worse than this. Using a tenth of the data, and in only two epochs, transfer learning beats it easily. This should be your default for images.

import torch
from torch import nn
from torchvision import datasets, models

weights = models.ResNet18_Weights.DEFAULT
tf = weights.transforms()
train = datasets.CIFAR10("data", train=True, download=True, transform=tf)
test = datasets.CIFAR10("data", train=False, download=True, transform=tf)
train = torch.utils.data.Subset(train, range(1200))
test = torch.utils.data.Subset(test, range(500))

model = models.resnet18(weights=weights)
for p in model.parameters():
    p.requires_grad = False
model.fc = nn.Linear(512, 10)

opt = torch.optim.AdamW(model.fc.parameters(), lr=1e-3)
dl = torch.utils.data.DataLoader(train, batch_size=64, shuffle=True)
tl = torch.utils.data.DataLoader(test, batch_size=128)

for epoch in range(2):
    model.train()
    for x, y in dl:
        loss = nn.functional.cross_entropy(model(x), y)
        opt.zero_grad(); loss.backward(); opt.step()

model.eval()
hits = 0
with torch.no_grad():
    for x, y in tl:
        hits += (model(x).argmax(1) == y).sum().item()
print(f"Images used to train: {len(train)}")
print(f"Accuracy: {hits / len(test):.3f}")

Images used to train: 1200
Accuracy: 0.482

Summary