Skip to main content

Module 9 — Transfer learning and progressive fine-tuning

Nine hundred thousand images and a week of GPU are what it took to train ResNet50 on ImageNet from scratch. Nobody does that anymore. Instead, you download the pretrained weights, replace the classifier head, and fine-tune on your task with the three thousand waste photographs from module 8. Done right, this reaches 90 % accuracy in an hour on a laptop with a modest GPU. Done wrong — and there is one very specific wrong way — the accuracy collapses to random and no error is reported. This module walks through the two-phase procedure and stops at every one of its traps.

Why transfer works

The first few convolutional layers of any CNN trained on natural images learn edge detectors, colour blobs and simple textures — the Sobel-like filters we met in module 1. Those features are not "cat features" or "waste features"; they are image features. They transfer almost losslessly from one visual dataset to another.

The middle layers combine those primitives into parts: circles, corners, fur patterns, glass reflections. Some of those transfer; some are ImageNet-specific.

The last few layers, and especially the classifier head, are strongly specialised to the source task. They must be replaced.

Phase 1: feature extraction with a frozen backbone

The safe first phase treats ResNet50 as a fixed feature extractor. Freeze every weight, replace the head, train only the head on the new task. The whole procedure fits in one screen.

import tensorflow as tf
from tensorflow.keras import layers

IMG_SIZE = 224
NUM_CLASSES = 6 # waste dataset

inputs = tf.keras.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
x = tf.keras.applications.resnet50.preprocess_input(inputs)

backbone = tf.keras.applications.ResNet50(
weights="imagenet", include_top=False, input_shape=(IMG_SIZE, IMG_SIZE, 3)
)
backbone.trainable = False # freeze the whole body

x = backbone(x, training=False) # see BatchNorm trap below
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dropout(0.3)(x)
outputs = layers.Dense(NUM_CLASSES, activation="softmax")(x)
model = tf.keras.Model(inputs, outputs)

model.compile(
optimizer=tf.keras.optimizers.Adam(1e-3),
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
)

On the waste dataset, five to ten epochs at learning rate 10310^{-3} take the head from random to about 85 % validation accuracy. Because only the head is trainable — a few thousand parameters instead of 25 million — the training is fast and the risk of overfitting is small.

The BatchNorm trap: trainable=False is not enough

Here is the trap that costs the most silently. backbone.trainable = False freezes the learnable parameters of every batch normalisation layer (the gamma and beta scale/shift). It does not switch the layer into inference mode. Inside a BN layer, two paths exist:

  • In training mode: compute batch statistics from the current batch, use them to normalise, and update running mean/variance.
  • In inference mode: use the stored running mean/variance, and do not update them.

If you leave BN in training mode on a frozen backbone, the ImageNet running statistics get overwritten by batch statistics from your six-class waste dataset, whose distribution is very different. Every subsequent forward pass then sees a normalisation trained on the wrong data. Symptom: your model looks fine on train but collapses on validation.

The fix is to pass training=False explicitly at the call site, as in the code above. That single argument keeps BN in inference mode regardless of the outer training state.

backbone.trainable = False freezes weights, not modes

The mode switch of BatchNorm is orthogonal to the weight freeze. The correct incantation is:

  • freeze the weights: backbone.trainable = False
  • keep BN in inference: backbone(x, training=False) at the call site. Both, not one or the other. This is the most frequent bug in Keras transfer learning, and it never raises an exception.

Phase 2: staged unfreezing with a differentiated learning rate

Once the head has stabilised, we cautiously unfreeze the top of the backbone. Do not unfreeze everything at once with the same learning rate you used for the head: gradients at 10310^{-3} will destroy the pretrained features in one epoch.

The recipe:

  1. Unfreeze the top N layers of the backbone; keep the bottom ones frozen.
  2. Recompile with a much smaller learning rate, typically 10× smaller than phase 1.
  3. Train for a few more epochs.
# Phase 2: unfreeze the last third of the ResNet50.
backbone.trainable = True
for layer in backbone.layers[:100]: # ~ first third of ResNet50
layer.trainable = False

# Keep BN layers frozen even in unfrozen blocks: they were trained on ImageNet
# statistics that the small waste dataset cannot reliably re-estimate.
for layer in backbone.layers:
if isinstance(layer, tf.keras.layers.BatchNormalization):
layer.trainable = False

model.compile(
optimizer=tf.keras.optimizers.Adam(1e-5),
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
)

A discriminative learning rate is the natural next step. Lower layers get an even smaller rate than upper layers: the deeper into the backbone, the more specialised the features and the less you want to disturb them. Libraries like fastai build this into their trainer; in Keras it is a few lines of custom optimiser configuration.

On the waste dataset, this second phase adds five points of accuracy over the frozen-backbone version and takes it above 90 %.

When transfer fails

Transfer learning is not magic. Three situations degrade it, sometimes to the point of being worse than training from scratch.

  • Domain gap too large. ImageNet is natural photographs; your task is satellite hyperspectral imagery, or medical CT slices, or industrial X-ray of welded steel. The low-level features (edges, colour blobs) transfer, but everything above is misleading. Symptom: transfer plateaus at a modest accuracy while a from-scratch model, given more data or more time, eventually beats it.
  • Input resolution mismatch. If you train the head at 96 by 96 to save memory, then fine-tune at 224 by 224, the receptive fields are effectively different scales. Fix: use the same resolution in both phases, and prefer the one the backbone was pretrained at (224 for ResNet50).
  • Excessive freezing on a large target dataset. Transfer is a data-efficiency trick; with a million target-domain images and time to train, you can and should unfreeze everything. The frozen backbone was appropriate for small datasets; on a large one, it becomes a bottleneck.
Save every phase

Save a checkpoint at the end of phase 1 before you touch phase 2. If phase 2 destroys the model (it will, at least once), phase 1 is a good baseline to roll back to. Checkpoints are cheap, model outages are expensive.

PyTorch equivalent

In PyTorch, freezing is done by setting param.requires_grad = False on the parameters you want to freeze, and BatchNorm's inference mode is toggled by calling model.eval() on the frozen submodule (or by inheriting from torch.nn.BatchNorm2d and overriding train()). Same trap, same fix, different API.

Bringing it back to the red thread

At this point in the course, the waste-sorting model is around 91 % validation accuracy — but two questions remain. Is that number honest? And when the model is wrong, what did it actually look at? Module 10 answers both with Grad-CAM.

In summary

  • Transfer learning splits training into two phases: feature extraction on a frozen backbone with the head only, then fine-tuning with the top of the backbone unfrozen and a much smaller learning rate.
  • backbone.trainable = False freezes the weights, not the BatchNorm mode. Pass training=False at the call site to keep BN in inference mode; skipping this is the most frequent silent failure in Keras transfer.
  • Discriminative learning rates give deeper layers even smaller steps than upper ones, preserving what was learned from ImageNet.
  • Transfer fails on large domain gaps, on resolution mismatches and on large target datasets where full training would eventually do better. Save a phase-1 checkpoint before every phase-2 attempt.

Next module: opening the model. Grad-CAM shows where the fine-tuned ResNet50 is actually looking when it decides "cardboard" versus "paper", and reveals when it is looking at the background instead of the object.