How to Use Transfer Learning in PyTorch With a Pretrained ResNet Model

Two models, same architecture, same data, same settings, trained for the same three epochs. One reached 97.78% accuracy. The other reached 33.33%, which on a three way choice is precisely what you would get by flipping a coin with three sides. The only difference between them was where their starting numbers came from. That gap is the single most useful thing a beginner can learn about deep learning, because it means the hardest part of training a good model has usually already been done by somebody else, on hardware you do not have, using data you could never collect. This tutorial shows you how to pick that work up and use it.

What Transfer Learning Actually Transfers

A network trained on a large image collection does not learn one skill. It learns a stack of them, and they are not equally specific to the task it was trained on.

The earliest layers become detectors for edges, corners, and colour transitions. Slightly deeper layers combine those into textures and simple repeating patterns. Deeper still, they respond to shapes and object parts. Only the final layers become genuinely specific to the exact categories the network was asked about.

Edges and textures are not properties of any particular dataset. They are properties of images. A network that detects them well is useful for medical scans, manufacturing defects, and satellite tiles, none of which resembled its training data.

Why This Matters More Than Any Architecture Choice

Beginners tend to assume that better results come from picking a better model. Far more often they come from starting a model in a better place.

The comparison run for this article makes the point bluntly. Two identical ResNet18 networks, the same 30 training images, the same optimizer, the same three epochs. The one starting from pretrained weights reached 97.78%. The one starting from random weights reached 33.33%, which on a three way choice is exactly chance. The architecture was not the variable. The starting point was.

Feature Extraction Compared With Fine Tuning

There are two ways to build on a pretrained model, and choosing between them is mostly a question of how much data you have.

Feature extraction freezes the entire pretrained network and trains only a new final layer. The frozen part becomes a fixed function that converts an image into a list of numbers describing what is in it, and the new layer learns to map those numbers onto your categories. This is what the tutorial below does.

Fine tuning additionally unfreezes some of the pretrained layers so they can adapt. It reaches higher accuracy when you have enough data, and it requires a much smaller learning rate on the unfrozen layers, because large updates will destroy the features you were trying to build on.

Step 1: Prepare Data the Way the Model Expects

This step is where silent accuracy loss happens. A pretrained model expects input formatted exactly as it was during its original training.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms, models

torch.manual_seed(42)

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

tfm = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])

full_train = datasets.ImageFolder("shapes/train", transform=tfm)
val_set    = datasets.ImageFolder("shapes/val",   transform=tfm)

Two constraints are non negotiable.

Images must be 224 by 224. That is the resolution ResNet was trained on. Other sizes technically pass through because of adaptive pooling, but the learned features are tuned to this scale and accuracy suffers at very different resolutions.

Normalisation must use ImageNet statistics. Those six numbers are the per channel mean and standard deviation of the ImageNet training set. The pretrained weights expect input centred that way. Substituting your own dataset statistics, or skipping normalisation entirely, produces no error and quietly costs you accuracy.

Now restrict training to 10 images per class, which is what makes the comparison meaningful:

idx, per_class = [], {0: 0, 1: 0, 2: 0}
for i, (_, lbl) in enumerate(full_train.samples):
    if per_class[lbl] < 10:
        idx.append(i)
        per_class[lbl] += 1
train_set = Subset(full_train, idx)

print("Classes           :", full_train.classes)
print("Training images   :", len(train_set), "(10 per class)")
print("Validation images :", len(val_set))

Classes           : ['circle', 'square', 'triangle']
Training images   : 30 (10 per class)
Validation images : 90

Thirty training images is a deliberately punishing amount. It is also realistic, because scarce labelled data is the normal situation outside of benchmarks and the exact circumstance transfer learning exists to address.

Step 2: Load the Pretrained Model and Freeze It

Three operations: load with weights, freeze everything, replace the head.

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

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

model.fc = nn.Linear(model.fc.in_features, 3)

print("Original fc layer:", models.resnet18().fc)
print("Replaced fc layer:", model.fc)

total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total parameters    : {total:,}")
print(f"Trainable parameters: {trainable:,}  ({100*trainable/total:.2f}%)")

Original fc layer: Linear(in_features=512, out_features=1000, bias=True)
Replaced fc layer: Linear(in_features=512, out_features=3, bias=True)
Total parameters    : 11,178,051
Trainable parameters: 1,539  (0.01%)

That final line is the whole technique in one number. You are training 1,539 parameters out of 11.2 million.

The count of 1,539 is arithmetic you can verify: 512 input features times 3 output classes gives 1,536 weights, plus 3 bias terms.

Step 3: Train Only the Head

Pass the optimizer just the parameters that require gradients.

train_loader = DataLoader(train_set, batch_size=10, shuffle=True)
val_loader   = DataLoader(val_set,   batch_size=30)

params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.Adam(params, lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

def accuracy(model, loader):
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for X, y in loader:
            correct += (model(X).argmax(1) == y).sum().item()
            total += y.size(0)
    return 100 * correct / total

print(f"Accuracy before training: {accuracy(model, val_loader):.2f}%")

for epoch in range(1, 4):
    model.train()
    running = 0.0
    for X, y in train_loader:
        loss = loss_fn(model(X), y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running += loss.item()
    print(f"epoch {epoch} | loss {running/len(train_loader):.4f} | "
          f"val accuracy {accuracy(model, val_loader):6.2f}%")

Accuracy before training: 16.67%
epoch 1 | loss 1.3806 | val accuracy  65.56%
epoch 2 | loss 0.9656 | val accuracy  88.89%
epoch 3 | loss 0.7533 | val accuracy  97.78%

From 16.67% to 97.78% on 30 training images.

Filtering the parameter list is a small but real optimisation. Adam stores two additional tensors of optimizer state per parameter it manages, so handing it all 11.2 million would allocate memory for state on weights that never change.

Step 4: Run the Same Model From Scratch

The comparison is what makes the result credible. Identical architecture, identical data, identical hyperparameters, random initial weights.

scratch = models.resnet18(weights=None)
scratch.fc = nn.Linear(scratch.fc.in_features, 3)
optimizer = torch.optim.Adam(scratch.parameters(), lr=1e-3)

Trainable parameters: 11,178,051
Accuracy before training: 33.33%
epoch 1 | loss 1.5395 | val accuracy  31.11%
epoch 2 | loss 1.1269 | val accuracy  37.78%
epoch 3 | loss 0.8267 | val accuracy  33.33%

Loss falls from 1.5395 to 0.8267, so the model is genuinely learning something. Validation accuracy stays at chance level, which means what it learned was the 30 specific training images rather than any property of circles, squares, and triangles.

The speed difference deserves a note, because it is the opposite of what people expect. Both models run the same forward pass, but the frozen model computes gradients for 1,539 parameters instead of 11.2 million, so the backward pass is dramatically cheaper. Absolute timings depend on your machine, which is why the table gives a ratio. Across two runs here the from scratch model took about one and a half times as long per epoch.

Step 5: Fine Tune by Unfreezing the Last Block

What you just did is feature extraction. Fine tuning unfreezes part of the backbone as well.

model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
for param in model.parameters():
    param.requires_grad = False

# unfreeze the final convolutional block
for param in model.layer4.parameters():
    param.requires_grad = True

model.fc = nn.Linear(model.fc.in_features, 3)

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

The two learning rates are the important detail. The new head starts from random weights and needs to move quickly. The unfrozen convolutional block already holds useful features, so it needs small adjustments rather than large ones. A single learning rate applied to both would destroy what layer4 already knows.

Start at the top row regardless of which one describes you. It trains in seconds and gives you a baseline, and you only need the rows below it if that baseline is not good enough.

Break It On Purpose

Trainable parameters: 0
ValueError: optimizer got an empty parameter list

Error Decoder

Choosing a Backbone

ResNet18 is a sensible default. It is small, fast, and well understood. Alternatives are worth knowing about.

ResNet50 has roughly twice the depth and usually a few points more accuracy, at proportionally higher compute cost. Use it when ResNet18 is close but not quite good enough.

EfficientNet models generally achieve better accuracy per parameter, which matters when deploying to constrained hardware.

Vision Transformers frequently outperform convolutional networks given large datasets, but they typically need more data to fine tune well, so they are a weaker choice for the small dataset case this article addresses.

Whichever you pick, the pattern is identical: load with weights, freeze, replace the final layer, train. Only the attribute name changes, since some architectures call the final layer classifier rather than fc. Print the model to check.

Conclusion

Transfer learning is the highest return technique available to anyone working with limited labelled data, and the measurements in this article show why. A frozen ResNet18 training 1,539 parameters reached 97.78% on 30 images. The identical architecture trained from scratch reached 33.33%, which is chance.

The recipe is four lines and rarely changes. Load the model with its pretrained weights, freeze the parameters, replace the final layer with one matching your class count, and train only what remains unfrozen. Normalise with the statistics the model was pretrained on, or you lose accuracy for no visible reason.

Start with everything frozen every time. It trains in seconds, needs almost no data, and gives you a number to beat. Unfreeze more layers only when that baseline falls short and you have the data to support it. Most projects never need to go further than the frozen version, which is a considerably better position than the one you are in when starting from random weights.

FAQS