Skip to main content

Module 8 — Transfer learning with pretrained models

Training a vision network from scratch takes millions of images and days of compute. Transfer learning gets better results with a few thousand images and a few minutes. This module explains why that works, then how to do it without falling into the trap that ruins half of all attempts.

Why early layers are reusable

A deep network does not build one representation: it stacks several, from the most general to the most specific.

Five-stage representation hierarchy of a vision network. Raw pixels feed early layers that detect edges and colour gradients, then middle layers that recognise textures and object parts, then late layers that identify whole objects, and finally a classification head specific to the original dataset.Raw pixelscolour valuesno structureEarly layersedges, cornerscolour gradientsuniversalMiddle layerstextures, repeated motifsobject partsLate layerswhole objectsspecific to classes seenClassification headprobabilities over 1000 classesto be replaced
The deeper you go, the more the representation is specific to the original dataset. The first two stages detect edges useful to any visual task at all: that generality is what makes transfer possible. Only the head must always be replaced.

An edge detector remains an edge detector whether the image shows a cat or a chest radiograph. That is what justifies keeping the weights of a network trained on ImageNet and replacing only its final part.

Loading a model without its head

from tensorflow import keras

base = keras.applications.EfficientNetB0(
weights="imagenet",
include_top=False, # drops the 1000-class head
input_shape=(224, 224, 3),
)
base.trainable = False # freezes the whole base

include_top=False removes the original classification layer, useless since your classes are not ImageNet's. trainable = False freezes the weights: they will no longer receive a gradient.

You then add your own head, using the functional API from module 3:

from tensorflow.keras import layers

inputs = keras.Input(shape=(224, 224, 3))
x = keras.applications.efficientnet.preprocess_input(inputs)
x = base(x, training=False)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dropout(0.3)(x)
outputs = layers.Dense(num_classes, activation="softmax")(x)

model = keras.Model(inputs, outputs)

The preprocess_input specific to each model family is not optional. Every architecture expects its inputs in a particular range — some in [0,1][0, 1], others in [1,1][-1, 1], others centred on ImageNet means. Feeding raw 0-to-255 pixels to a network that expects [1,1][-1, 1] gives mediocre results with no error message whatsoever.

The BatchNormalization trap

Here is the mistake behind most failed transfers, and it is almost invisible.

A BatchNormalization layer holds two kinds of quantity: learned weights, and running statistics — mean and variance accumulated during the original training. Now, base.trainable = False freezes the weights, but it is not enough to freeze the statistics: if the layer is called in training mode, it keeps updating them with your data.

On a small dataset those statistics drift, move away from the ones the frozen weights were optimised for, and destroy the very representation you were trying to preserve. The symptom is accuracy that stalls or regresses while everything looks correct.

The fix is the training=False in the base(x, training=False) call above. It forces inference mode for the entire base, statistics included.

Two distinct mechanisms, frequently conflated

base.trainable = False decides which weights receive a gradient. The training=False argument decides how the layers behave. Both are necessary and neither replaces the other. On a base containing batch normalisation — which is nearly every modern architecture — omitting the second cancels most of the first's benefit.

Two regimes, in this order

The effective approach runs in two phases, and the order is not negotiable.

Phase 1, feature extraction. The base is frozen, only the new head trains, with an ordinary learning rate.

model.compile(
optimizer=keras.optimizers.Adam(1e-3),
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
)
model.fit(dataset, validation_data=val_dataset, epochs=10)

Phase 2, fine-tuning. You unfreeze part of the base and resume training with a much lower rate.

base.trainable = True
for layer in base.layers[:-30]: # unfreeze only the last 30
layer.trainable = False

model.compile(
optimizer=keras.optimizers.Adam(1e-5), # a hundred times lower
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
)
model.fit(dataset, validation_data=val_dataset, epochs=10)

Recompiling after changing trainable is mandatory: without it, the list of trainable variables stays the one established at the previous compile and the unfreezing has no effect.

The very low rate has a simple explanation. The new head, randomly initialised, produces enormous gradients at first. Applied to a pretrained base at an ordinary rate, they erase within a few batches the information accumulated over millions of images. It is also why phase 1 must precede phase 2: it brings the head to a reasonable state before the base is exposed.

How much to unfreeze

SituationStrategy
little data, similar domainfreeze everything, train the head alone
plenty of data, similar domainunfreeze the last third
little data, distant domainunfreeze middle layers, keep early ones frozen
plenty of data, very distant domainconsider training from scratch

The "little data, distant domain" case is the trickiest, and medical imaging is the textbook example. Early layers stay useful, late ones are too specific to ImageNet objects, and there is not enough data to relearn everything. Unfreezing from the middle is the compromise, to be validated experimentally.

Key takeaways

  • Transfer works because early layers learn universal patterns — edges, textures — independent of the original task.
  • include_top=False drops the original head, and the family's preprocess_input is essential, since each architecture expects inputs in a specific range.
  • trainable = False freezes weights but not batch normalisation statistics: without training=False at the call site, they drift and destroy the representation you preserved.
  • Always feature extraction first, fine-tuning second at a hundred-fold lower rate, and recompile after every change to trainable.

Next module: distributed training, for when a single accelerator is no longer enough.