🎓 A1Learner
🔍

PyTorch What Transfer Learning Is

Transfer learning reuses a trained model's early layers on a new task instead of training from zero. The reused layers are called a feature extractor, because they turn raw input into useful features.

Two tasks, two full models

A second task trained from zero pays for a whole second network. Every layer is learned again, even the early ones.

separate = 2 * sum(p.numel() for p in model.parameters())
print("Two separate models:", separate, "parameters")

One extractor, two small heads

Transfer learning trains the shared part once. Each new task then adds only a small head.

shared = sum(p.numel() for p in model.features.parameters())
shared += 2 * sum(p.numel() for p in model.classifier.parameters())
print("One extractor, two heads:", shared, "parameters")

The extractor output is small

The extractor squeezes one image into 16 numbers. A head only has to read those 16 numbers.

one = torch.zeros(1, 1, 16, 16)
feats = torch.flatten(model.features(one), 1)
print("Feature shape for one image:", tuple(feats.shape))

Example

import torch
from torch import nn

torch.manual_seed(0)


class ShapesCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 8, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(8, 16, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Linear(16, 3)

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)


model = ShapesCNN()

print("The feature extractor:")
for layer in model.features:
    print(" ", layer)

total = sum(p.numel() for p in model.features.parameters())
print("Feature extractor parameters:", total)
print("Classifier parameters:", sum(p.numel() for p in model.classifier.parameters()))

Key points

  • Feature extractor: the early layers of a trained model, kept for reuse on new tasks
  • Head: the small final layer that maps features to the new task's classes
  • Use transfer learning when the new task has few labelled rows
  • This course pretrains on synthetic shapes data, so no download is needed

In this example, the shapes CNN from the CNN project is built with its seed set. The model splits into two parts. The features part holds the two conv layers, the ReLU layers, and the pooling layers. The classifier part is one Linear layer. The program prints each extractor layer in order. Then it counts the parameters on each side. The extractor holds 1,248 parameters. The classifier holds 51. The extractor is the reusable part, and the next lessons put it to work.

Note

Real projects often start from models pretrained on large public datasets, with the same mechanism one shown here.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What is a feature extractor?