XenorsTransfer Learning: Pre-trained Models in PyTorch and Keras
Training a CNN from scratch needs large data and compute. Transfer learning starts from a model trained on ImageNet and adapts it to your task with far less of both.
Feature extraction versus fine-tuning
Feature extraction freezes the whole backbone and trains only a new head. Fine-tuning also updates some or all backbone layers. Use feature extraction first when data is small, and fine-tune when you have more data or your images differ a lot from ImageNet.
| Situation | Strategy |
|---|---|
| Small data, similar domain | Freeze backbone, train head |
| Medium data, similar domain | Fine-tune top layers |
| Large data, different domain | Fine-tune most or all layers |
Transfer learning in PyTorch
Replace the classifier, freeze everything else, and give the optimiser only the parameters that train.
import torch, torch.nn as nn
from torchvision import models
weights = models.ResNet50_Weights.DEFAULT
model = models.resnet50(weights=weights)
for p in model.parameters():
p.requires_grad = False
model.fc = nn.Linear(model.fc.in_features, 5) # 5 classes
opt = torch.optim.Adam(model.fc.parameters(), lr=1e-3)
prep = weights.transforms() # use these for your data
# Phase 2: unfreeze the last block and lower the learning rate
for p in model.layer4.parameters():
p.requires_grad = True
opt = torch.optim.Adam([
{"params": model.layer4.parameters(), "lr": 1e-5},
{"params": model.fc.parameters(), "lr": 1e-4}])Transfer learning in Keras
Keras applications include preprocessing helpers. Keep batch-norm layers in inference mode while the base is frozen.
import tensorflow as tf
from tensorflow import keras
base = keras.applications.ResNet50(weights="imagenet", include_top=False,
input_shape=(224, 224, 3))
base.trainable = False
inputs = keras.Input((224, 224, 3))
x = keras.applications.resnet50.preprocess_input(inputs)
x = base(x, training=False)
x = keras.layers.GlobalAveragePooling2D()(x)
x = keras.layers.Dropout(0.2)(x)
outputs = keras.layers.Dense(5, activation="softmax")(x)
model = keras.Model(inputs, outputs)
model.compile("adam", "sparse_categorical_crossentropy", metrics=["accuracy"])
# model.fit(train_ds, validation_data=val_ds, epochs=5)
# Fine-tune: unfreeze the top of the base with a tiny learning rate
base.trainable = True
for layer in base.layers[:-30]:
layer.trainable = False
model.compile(keras.optimizers.Adam(1e-5),
"sparse_categorical_crossentropy", metrics=["accuracy"])Avoid common mistakes
Do not skip the model's own preprocessing. Do not fine-tune with the same learning rate you used for the head, because large updates destroy useful pre-trained features. Always hold out a validation set and watch for overfitting when the backbone is unfrozen.
Hands-on practice
Run the examples in the downloadable Jupyter notebook. The extra checks below focus on mistakes that appear in real projects, not just syntax.
Freeze a pretrained backbone
from torchvision.models import resnet18, ResNet18_Weights
model=resnet18(weights=ResNet18_Weights.DEFAULT)
for p in model.parameters(): p.requires_grad=False
print(sum(p.requires_grad for p in model.parameters()))0
Start with frozen features on small datasets, then unfreeze gradually if validation improves.
Replace the classifier head
import torch.nn as nn
model.fc=nn.Linear(model.fc.in_features,3)
print(model.fc)Linear(in_features=512, out_features=3, bias=True)
The new head is trainable even when the backbone remains frozen.
Self-check
- Can you explain why the first practice result is trustworthy rather than just plausible?
- What would you monitor or test before using this Transfer Learning workflow in production?
- Which assumption in the example would break first on a larger or messier dataset?
Frequently asked questions
When should I use transfer learning?
Almost always for image tasks with limited data. Starting from pre-trained weights usually beats training from scratch.
Why use a lower learning rate for fine-tuning?
Pre-trained weights already encode useful features. A small learning rate adjusts them gently instead of overwriting them.
PyTorch or Keras for transfer learning?
Both work well. Keras is faster to prototype, while PyTorch gives finer control of the training loop.
Primary references
Use these official sources when exact syntax, defaults or definitions matter beyond the examples here.
