9.5. Transfer Learning with Pretrained Models

What Transfer Learning Is in Practice

Transfer learning means starting from a model that has already learned useful visual features on a large dataset, then reusing those learned weights for your task. In computer vision, this usually means taking a convolutional neural network trained on ImageNet and adapting it to a new dataset with fewer labeled images. Early layers often learn generic features like edges and textures, while later layers become more task specific. Transfer learning takes advantage of that, so you train faster, need less data, and often get better accuracy than training from scratch.

When to Use Transfer Learning

Transfer learning is most helpful when your dataset is small to medium sized, or when you want a strong baseline quickly. It is also useful when your images are similar to the domain the model was pretrained on, which for ImageNet typically means natural photos. If your images are very different, such as medical scans or satellite imagery, transfer learning can still help, but you may need more fine tuning and you should expect smaller gains.

Choosing a Pretrained Model in PyTorch

In PyTorch, pretrained vision models are commonly accessed through torchvision.models. Modern torchvision versions use a weights= argument rather than pretrained=True. You typically pick a family like ResNet, EfficientNet, MobileNet, or ViT depending on your constraints. Larger models often give better accuracy but require more memory and compute. Smaller models train and run faster and are easier to deploy.

A typical workflow is to load a model with pretrained weights, replace the final classification layer to match your number of classes, then train.

Loading a Pretrained Model and Inspecting the Classifier Head

Most pretrained models are built as nn.Module objects, so you can inspect their final layer to see where to plug in your new head. For example, many ResNet variants expose the final classifier as model.fc, while MobileNet and EfficientNet often use model.classifier.

Here is a common pattern with ResNet:

import torch
import torch.nn as nn
from torchvision import models
num_classes = 5
weights = models.ResNet18_Weights.DEFAULT
model = models.resnet18(weights=weights)
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, num_classes)

The key idea is that you keep the pretrained feature extractor, and you swap the last layer so its output shape is [batch_size, num_classes].

Matching the Pretrained Model Preprocessing

Pretrained weights come with an expected preprocessing recipe, typically including resizing, center cropping, conversion to tensor, and normalization with specific mean and standard deviation. In torchvision, the weights object provides the right transforms.

from torchvision import transforms
weights = models.ResNet18_Weights.DEFAULT
preprocess = weights.transforms()

You can then apply preprocess inside your dataset pipeline. This is important because the model was trained assuming inputs were normalized in a particular way.

Two Main Strategies: Feature Extraction vs Fine Tuning

There are two standard ways to use a pretrained model. Feature extraction means you freeze the pretrained layers and only train the new head. Fine tuning means you train some or all of the pretrained layers along with the new head.

Feature extraction is faster and reduces the risk of overfitting when your dataset is small. Fine tuning can yield better performance when you have more data or your task differs from ImageNet.

The mechanics of freezing are simple, you set requires_grad = False for parameters you do not want to update:

for param in model.parameters():
    param.requires_grad = False
model.fc = nn.Linear(model.fc.in_features, num_classes)

Now only the new model.fc parameters will receive gradients and be updated by the optimizer.

Setting Up the Optimizer Correctly for Transfer Learning

When you freeze layers, you should pass only trainable parameters to the optimizer. This avoids wasted computation and makes your intent clear.

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

When you fine tune the whole network, you typically use a smaller learning rate than training from scratch, because you are adjusting already good weights rather than learning everything anew.

A Common Fine Tuning Pattern: Unfreeze Gradually

A practical approach is to start with feature extraction for a few epochs, then unfreeze part of the backbone and continue training with a lower learning rate. This reduces the chance of unstable updates early on.

In code, this looks like: first freeze all backbone parameters, train the head, then unfreeze selected layers.

Exactly which layers to unfreeze depends on the architecture. With ResNet, you might unfreeze layer4 first, then more if needed.

Using Different Learning Rates for Backbone and Head

Often you want the new head to learn faster than the pretrained backbone. You can use optimizer parameter groups to assign different learning rates.

backbone_params = []
head_params = []
for name, param in model.named_parameters():
    if "fc" in name:
        head_params.append(param)
    else:
        backbone_params.append(param)
optimizer = torch.optim.SGD(
    [
        {"params": backbone_params, "lr": 1e-4},
        {"params": head_params, "lr": 1e-3},
    ],
    momentum=0.9
)

This setup is especially useful when you unfreeze the backbone for fine tuning.

Handling Different Numbers of Input Channels

Pretrained models typically expect 3 channel RGB images. If your data is grayscale, you have three common options. You can convert grayscale to 3 channels by repeating the channel in your transform. You can modify the first convolution layer to accept 1 channel and initialize it sensibly. Or you can train from scratch if neither is appropriate.

The simplest beginner friendly approach is to convert to 3 channels in preprocessing, so the network architecture stays unchanged.

Transfer Learning for Feature Extraction Without Training the Backbone

Sometimes you want to use a pretrained CNN as a fixed feature extractor and train a separate classifier outside the network, for example a linear classifier or a small MLP. In PyTorch, you can do this by removing the final layer and using the intermediate embedding. Many models can be turned into feature extractors by replacing the classifier with nn.Identity().

model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
model.fc = nn.Identity()
model.eval()

Now the output is a feature vector per image, which you can feed to another model. This approach is also helpful for quick experiments and debugging.

Common Pitfalls Specific to Transfer Learning

A frequent issue is forgetting to switch preprocessing to match the pretrained weights, which can make training appear to stall. Another is replacing the wrong layer, which leads to shape mismatches or a model that still outputs 1000 classes. Also, if you freeze everything, including the new head by accident, the loss will not improve because no parameters are being updated.

Batch normalization layers deserve attention. If you freeze the backbone, you often still want stable behavior from batch norm. In many beginner setups, keeping the model in training mode while freezing parameters works, but batch norm running statistics may still update. If your batch size is very small, this can introduce noise. A simple mitigation is to keep the backbone in eval mode while training only the head, then switch back when you start fine tuning.

What You Should Take Away

Transfer learning in PyTorch boils down to loading pretrained weights, applying the correct transforms, replacing the classifier head, and choosing whether to freeze or fine tune. Once those pieces are correct, you can reuse your standard training loop, loss, and metrics from earlier chapters, and you will usually get a strong image model with much less training effort.

Views: 86