PyTorch Transfer Learning

Transfer learning in PyTorch means starting from a model already trained on a huge dataset, such as ImageNet, and adapting it to your own task. The early layers have learned general features like edges and textures, so you only need a small dataset and short training to get excellent results.

Loading a Pretrained Model

torchvision uses weights enums. The old pretrained=True argument has been removed.

import torch
from torch import nn
from torchvision import models

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

Linear(in_features=512, out_features=1000, bias=True)

The final fc layer predicts the 1000 ImageNet classes. We will replace it.

Freeze the Backbone and Replace the Head

for p in model.parameters():
    p.requires_grad = False

num_classes = 2   # e.g. cats vs dogs
model.fc = nn.Linear(model.fc.in_features, num_classes)

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(trainable, total)

1026 11177538

The new layer is created after freezing, so it is trainable by default. Only 1,026 of about 11 million parameters will be updated, which makes training fast even on a CPU.

Preparing Your Images

Use the same preprocessing the model was trained with. The weights object provides it. ImageFolder reads a directory where each sub-folder is one class:

from torchvision import datasets
from torch.utils.data import DataLoader

preprocess = weights.transforms()
train_ds = datasets.ImageFolder("pets/train", transform=preprocess)   # pets/train/cat, pets/train/dog
train_dl = DataLoader(train_ds, batch_size=32, shuffle=True)
print(train_ds.classes)

['cat', 'dog']

Training the Head

opt = torch.optim.AdamW(model.fc.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(3):
    model.train()
    for xb, yb in train_dl:
        loss = loss_fn(model(xb), yb)
        opt.zero_grad()
        loss.backward()
        opt.step()
    print(f"epoch {epoch + 1}: loss {loss.item():.3f}")

epoch 1: loss 0.284
epoch 2: loss 0.151
epoch 3: loss 0.097
(sample output)

Fine-Tuning Deeper Layers

Once the head is trained, you can unfreeze the last block and continue with a smaller learning rate for a further boost:

for p in model.layer4.parameters():
    p.requires_grad = True

opt = torch.optim.AdamW([
    {"params": model.layer4.parameters(), "lr": 1e-4},
    {"params": model.fc.parameters(), "lr": 1e-3},
])

(optimizer with two parameter groups created)

For text, the same idea applies with pretrained transformers from Hugging Face; see the Hugging Face tutorial for fine-tuning language models.

Frequently asked questions

What is transfer learning in PyTorch?

Reusing a model pretrained on a large dataset and adapting it to a new task, usually by replacing the final layer and training it, sometimes also fine-tuning deeper layers.

What replaced pretrained=True in torchvision?

The weights argument, for example models.resnet18(weights=models.ResNet18_Weights.DEFAULT). Use weights=None for random initialization.

Should I freeze layers when fine-tuning?

With small datasets, freeze the backbone and train only the head first. With more data, unfreeze later layers or the whole model with a low learning rate.

How many images do I need for transfer learning?

Often a few hundred per class is enough for good results, because the pretrained features already capture general visual patterns.