Skip to content

Transfer Learning - Using Pre-trained Models

What you’ll learn

  • why you almost never need to train a CNN like ResNet or Xception from scratch
  • how to load an ImageNet-pretrained model with a single line of keras.applicationskeras.applications
  • the difference between feature extraction (freeze the base) and fine-tuning (unfreeze it later)
  • a full transfer-learning workflow: freeze → train the head → unfreeze → fine-tune at a low learning rate

Why train from scratch when you don’t have to

Building a GoogLeNet or ResNet from scratch and training it on ImageNet takes serious compute and millions of labeled images. Most of us don’t have either. The good news: keras.applicationskeras.applications ships dozens of architectures already trained on ImageNet, ready to download with one line of code:

Loading a pretrained ImageNet classifier
import tensorflow as tf
 
model = tf.keras.applications.resnet50.ResNet50(weights="imagenet")
Loading a pretrained ImageNet classifier
import tensorflow as tf
 
model = tf.keras.applications.resnet50.ResNet50(weights="imagenet")

That’s it — this downloads ResNet-50’s architecture and its ImageNet weights. Before you can use it, resize your images to what the model expects (224 × 224 for ResNet-50), then run them through that model’s own preprocess_input()preprocess_input() function, since every architecture expects pixels scaled a little differently:

Using it to classify an image
images_resized = tf.image.resize(images, [224, 224])
inputs = tf.keras.applications.resnet50.preprocess_input(images_resized * 255)
 
Y_proba = model.predict(inputs)
top_3 = tf.keras.applications.resnet50.decode_predictions(Y_proba, top=3)
Using it to classify an image
images_resized = tf.image.resize(images, [224, 224])
inputs = tf.keras.applications.resnet50.preprocess_input(images_resized * 255)
 
Y_proba = model.predict(inputs)
top_3 = tf.keras.applications.resnet50.decode_predictions(Y_proba, top=3)

decode_predictions()decode_predictions() turns the raw 1,000-class probability vector into human-readable labels like ("n03877845", "palace", 0.41)("n03877845", "palace", 0.41) — the class ID, its name, and the model’s confidence.

This is great when your classes are already inside ImageNet’s 1,000 categories. But what if you want to classify something ImageNet never saw — say, five specific flower species? That’s where transfer learning comes in.

Feature extraction vs. fine-tuning

diagram Transfer learning architecture mermaid
A frozen pretrained base extracts general features; only the new head is trained on your dataset at first.

A pretrained CNN’s early layers learned very general visual features — edges, textures, corners — that are useful for almost any image task, not just the 1,000 ImageNet classes. The later layers are more specialized toward those specific classes. So the recipe is:

  1. Feature extraction — throw away the model’s final classification layer, keep everything before it (the “base”), freeze its weights so they can’t change, and bolt on a brand-new head trained only on your data.
  2. Fine-tuning (optional, later) — once your new head has learned something reasonable, unfreeze the base too and keep training the whole model, using a much smaller learning rate so you nudge the pretrained weights instead of destroying them.

Freezing at first matters for two reasons: it’s dramatically cheaper (you’re only training the small new head), and it prevents large, random gradients from the untrained head from wrecking carefully learned pretrained weights before they’ve had a chance to settle.

A full example: classifying flowers with Xception

Let’s reuse a pretrained Xception model to classify photos of flowers into five species — a task Xception never saw during its original ImageNet training.

transfer_learning_flowers.py
import tensorflow as tf
import tensorflow_datasets as tfds
 
# 1. Load the data and split it ourselves (tf_flowers ships only a train split)
dataset, info = tfds.load("tf_flowers", as_supervised=True, with_info=True)
n_classes = info.features["label"].num_classes  # 5
 
test_set, valid_set, train_set = tfds.load(
    "tf_flowers",
    split=["train[:10%]", "train[10%:25%]", "train[25%:]"],
    as_supervised=True,
)
 
# 2. Resize + preprocess for Xception, batch, and prefetch
def preprocess(image, label):
    resized = tf.image.resize(image, [224, 224])
    return tf.keras.applications.xception.preprocess_input(resized), label
 
batch_size = 32
train_set = train_set.shuffle(1000).map(preprocess).batch(batch_size).prefetch(1)
valid_set = valid_set.map(preprocess).batch(batch_size).prefetch(1)
test_set = test_set.map(preprocess).batch(batch_size).prefetch(1)
 
# 3. Load Xception WITHOUT its final classification layers
base_model = tf.keras.applications.xception.Xception(
    weights="imagenet", include_top=False
)
avg = tf.keras.layers.GlobalAveragePooling2D()(base_model.output)
output = tf.keras.layers.Dense(n_classes, activation="softmax")(avg)
model = tf.keras.Model(inputs=base_model.input, outputs=output)
 
# 4. Freeze the base and train only the new head
for layer in base_model.layers:
    layer.trainable = False
 
optimizer = tf.keras.optimizers.SGD(learning_rate=0.2, momentum=0.9)
model.compile(loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])
model.fit(train_set, epochs=5, validation_data=valid_set)
 
# 5. Unfreeze and fine-tune the WHOLE model with a much lower learning rate
for layer in base_model.layers:
    layer.trainable = True
 
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.9)
model.compile(loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])
model.fit(train_set, epochs=10, validation_data=valid_set)
transfer_learning_flowers.py
import tensorflow as tf
import tensorflow_datasets as tfds
 
# 1. Load the data and split it ourselves (tf_flowers ships only a train split)
dataset, info = tfds.load("tf_flowers", as_supervised=True, with_info=True)
n_classes = info.features["label"].num_classes  # 5
 
test_set, valid_set, train_set = tfds.load(
    "tf_flowers",
    split=["train[:10%]", "train[10%:25%]", "train[25%:]"],
    as_supervised=True,
)
 
# 2. Resize + preprocess for Xception, batch, and prefetch
def preprocess(image, label):
    resized = tf.image.resize(image, [224, 224])
    return tf.keras.applications.xception.preprocess_input(resized), label
 
batch_size = 32
train_set = train_set.shuffle(1000).map(preprocess).batch(batch_size).prefetch(1)
valid_set = valid_set.map(preprocess).batch(batch_size).prefetch(1)
test_set = test_set.map(preprocess).batch(batch_size).prefetch(1)
 
# 3. Load Xception WITHOUT its final classification layers
base_model = tf.keras.applications.xception.Xception(
    weights="imagenet", include_top=False
)
avg = tf.keras.layers.GlobalAveragePooling2D()(base_model.output)
output = tf.keras.layers.Dense(n_classes, activation="softmax")(avg)
model = tf.keras.Model(inputs=base_model.input, outputs=output)
 
# 4. Freeze the base and train only the new head
for layer in base_model.layers:
    layer.trainable = False
 
optimizer = tf.keras.optimizers.SGD(learning_rate=0.2, momentum=0.9)
model.compile(loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])
model.fit(train_set, epochs=5, validation_data=valid_set)
 
# 5. Unfreeze and fine-tune the WHOLE model with a much lower learning rate
for layer in base_model.layers:
    layer.trainable = True
 
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.9)
model.compile(loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])
model.fit(train_set, epochs=10, validation_data=valid_set)

include_top=Falseinclude_top=False strips off the average-pooling + dense-softmax layers that Xception used for its original 1,000 ImageNet classes, leaving just the convolutional “feature extractor.” From there, we bolt on our own GlobalAveragePooling2DGlobalAveragePooling2D + Dense(n_classes, "softmax")Dense(n_classes, "softmax") head.

After a few epochs of frozen training, validation accuracy typically climbs to somewhere around 75–80% and plateaus — a sign the new head has learned what it can from fixed features. Unfreezing and fine-tuning at a low learning rate usually pushes accuracy up toward 90%+, since now the pretrained filters can adapt slightly to flower photos specifically.

A more surgical fine-tune: unfreezing only the top layers

The workflow above unfreezes the entire base for the fine-tuning step, which works, but it’s not the only option — and it’s not always the best one. Chollet’s version of fine-tuning is more surgical: instead of unfreezing everything, unfreeze only the last few layers of the convolutional base and leave the rest frozen.

Two reasons to prefer that:

  • Early layers learn generic features (edges, colors, simple textures) that are useful for almost any image task — there’s little to gain, and fast-decreasing returns, from fine-tuning them. The later layers encode more specialized features, which are exactly the ones that benefit from being nudged toward your specific dataset.
  • Fewer trainable parameters means less overfitting risk. A base like VGG16 has around 15 million parameters — unfreezing all of it to train on a dataset of a few thousand images is asking for trouble. Unfreezing just the last convolutional block keeps the trainable parameter count small enough to fine-tune safely.
fine_tune_top_layers_only.py
import tensorflow as tf
 
conv_base = tf.keras.applications.vgg16.VGG16(weights="imagenet", include_top=False)
 
# Step 1: freeze the whole base and train your new head first (as before) --
# fine-tuning only makes sense once that head is no longer randomly initialized.
 
# Step 2: unfreeze the base, then re-freeze everything except the last few layers
conv_base.trainable = True
for layer in conv_base.layers[:-4]:
    layer.trainable = False
 
# Step 3: recompile (Keras only picks up trainable changes at compile time)
# with a very low learning rate, so you nudge -- not overwrite -- the
# pretrained weights in the layers you did unfreeze.
model.compile(
    loss="binary_crossentropy",
    optimizer=tf.keras.optimizers.RMSprop(learning_rate=1e-5),
    metrics=["accuracy"],
)
 
callbacks = [
    tf.keras.callbacks.ModelCheckpoint(
        filepath="fine_tuning.keras", save_best_only=True, monitor="val_loss"
    ),
]
model.fit(train_dataset, epochs=30, validation_data=validation_dataset, callbacks=callbacks)
fine_tune_top_layers_only.py
import tensorflow as tf
 
conv_base = tf.keras.applications.vgg16.VGG16(weights="imagenet", include_top=False)
 
# Step 1: freeze the whole base and train your new head first (as before) --
# fine-tuning only makes sense once that head is no longer randomly initialized.
 
# Step 2: unfreeze the base, then re-freeze everything except the last few layers
conv_base.trainable = True
for layer in conv_base.layers[:-4]:
    layer.trainable = False
 
# Step 3: recompile (Keras only picks up trainable changes at compile time)
# with a very low learning rate, so you nudge -- not overwrite -- the
# pretrained weights in the layers you did unfreeze.
model.compile(
    loss="binary_crossentropy",
    optimizer=tf.keras.optimizers.RMSprop(learning_rate=1e-5),
    metrics=["accuracy"],
)
 
callbacks = [
    tf.keras.callbacks.ModelCheckpoint(
        filepath="fine_tuning.keras", save_best_only=True, monitor="val_loss"
    ),
]
model.fit(train_dataset, epochs=30, validation_data=validation_dataset, callbacks=callbacks)

Either approach — unfreeze everything at a low learning rate, or unfreeze just the last few layers — is a legitimate fine-tuning strategy. Start with the “unfreeze everything” version from the flowers example above since it’s simpler; reach for the more surgical top-layers-only version when your dataset is small and you’re seeing overfitting creep back in once the base becomes trainable.

Mini-checkpoint

Why start with the base model frozen instead of fine-tuning everything at once?

  • A freshly attached head has random weights, so its first gradients are large and noisy. If the pretrained base were trainable from the start, those noisy gradients could wreck weeks of ImageNet training in a few batches. Freezing first lets the head “catch up” safely.

🧪 Try It Yourself

Exercise 1 – Freeze a Pretrained Base

Exercise 2 – Add a New Classification Head

Exercise 3 – Unfreeze for Fine-Tuning

Two ways to do feature extraction

There’s actually a choice to make about how you run the frozen base over your data, and it comes with a real speed/flexibility trade-off:

  1. Fast feature extraction (no augmentation) — run the frozen convolutional base over every training image once, save its output (the extracted features) to a NumPy array, and train a small, plain DenseDense classifier on those cached features. Since the expensive convolutional base only ever runs once per image, an epoch over a few thousand images can finish in well under a second on CPU. The catch: because the features are precomputed, you can’t apply data augmentation — there’s no convolutional base left in the trainable pipeline to feed augmented images through.

    fast_feature_extraction.py
    import numpy as np
    import tensorflow as tf
     
    conv_base = tf.keras.applications.vgg16.VGG16(weights="imagenet", include_top=False)
     
    def get_features_and_labels(dataset):
        all_features, all_labels = [], []
        for images, labels in dataset:
            preprocessed = tf.keras.applications.vgg16.preprocess_input(images)
            all_features.append(conv_base.predict(preprocessed))
            all_labels.append(labels)
        return np.concatenate(all_features), np.concatenate(all_labels)
     
    train_features, train_labels = get_features_and_labels(train_dataset)
    # train_features.shape == (num_samples, 5, 5, 512) for a 5x5x512 VGG16 output
     
    inputs = tf.keras.Input(shape=(5, 5, 512))
    x = tf.keras.layers.Flatten()(inputs)
    x = tf.keras.layers.Dense(256)(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
    model = tf.keras.Model(inputs, outputs)
    model.compile(loss="binary_crossentropy", optimizer="rmsprop", metrics=["accuracy"])
    model.fit(train_features, train_labels, epochs=20)
    fast_feature_extraction.py
    import numpy as np
    import tensorflow as tf
     
    conv_base = tf.keras.applications.vgg16.VGG16(weights="imagenet", include_top=False)
     
    def get_features_and_labels(dataset):
        all_features, all_labels = [], []
        for images, labels in dataset:
            preprocessed = tf.keras.applications.vgg16.preprocess_input(images)
            all_features.append(conv_base.predict(preprocessed))
            all_labels.append(labels)
        return np.concatenate(all_features), np.concatenate(all_labels)
     
    train_features, train_labels = get_features_and_labels(train_dataset)
    # train_features.shape == (num_samples, 5, 5, 512) for a 5x5x512 VGG16 output
     
    inputs = tf.keras.Input(shape=(5, 5, 512))
    x = tf.keras.layers.Flatten()(inputs)
    x = tf.keras.layers.Dense(256)(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
    model = tf.keras.Model(inputs, outputs)
    model.compile(loss="binary_crossentropy", optimizer="rmsprop", metrics=["accuracy"])
    model.fit(train_features, train_labels, epochs=20)
  2. Feature extraction together with data augmentation — chain a data augmentation stage, the frozen conv_baseconv_base, and a new classifier into one end-to-end model, so every training image is re-augmented and re-run through the whole convolutional base on every epoch. This is much slower (the expensive base now runs on every image, every epoch, instead of once total), but it lets you fight overfitting with augmentation while you still benefit from pretrained features — the best option when your small dataset is prone to overfitting even with a frozen base.

    feature_extraction_with_augmentation.py
    import tensorflow as tf
     
    data_augmentation = tf.keras.Sequential([
        tf.keras.layers.RandomFlip("horizontal"),
        tf.keras.layers.RandomRotation(0.1),
        tf.keras.layers.RandomZoom(0.2),
    ])
     
    conv_base = tf.keras.applications.vgg16.VGG16(weights="imagenet", include_top=False)
    conv_base.trainable = False   # freeze BEFORE building the model that wraps it
     
    inputs = tf.keras.Input(shape=(180, 180, 3))
    x = data_augmentation(inputs)
    x = tf.keras.applications.vgg16.preprocess_input(x)
    x = conv_base(x)
    x = tf.keras.layers.Flatten()(x)
    x = tf.keras.layers.Dense(256)(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
    model = tf.keras.Model(inputs, outputs)
    model.compile(loss="binary_crossentropy", optimizer="rmsprop", metrics=["accuracy"])
    feature_extraction_with_augmentation.py
    import tensorflow as tf
     
    data_augmentation = tf.keras.Sequential([
        tf.keras.layers.RandomFlip("horizontal"),
        tf.keras.layers.RandomRotation(0.1),
        tf.keras.layers.RandomZoom(0.2),
    ])
     
    conv_base = tf.keras.applications.vgg16.VGG16(weights="imagenet", include_top=False)
    conv_base.trainable = False   # freeze BEFORE building the model that wraps it
     
    inputs = tf.keras.Input(shape=(180, 180, 3))
    x = data_augmentation(inputs)
    x = tf.keras.applications.vgg16.preprocess_input(x)
    x = conv_base(x)
    x = tf.keras.layers.Flatten()(x)
    x = tf.keras.layers.Dense(256)(x)
    x = tf.keras.layers.Dropout(0.5)(x)
    outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
    model = tf.keras.Model(inputs, outputs)
    model.compile(loss="binary_crossentropy", optimizer="rmsprop", metrics=["accuracy"])

    With the base frozen, only the new DenseDense layers actually train — but every image still passes through data_augmentationdata_augmentation and the full conv_baseconv_base on every step, since they’re now part of the same model.

Next

Continue to Data Augmentation for Small Datasets — the technique that makes feature extraction with augmentation (and training a convnet from scratch) far more resistant to overfitting.

If this helped you, consider buying me a coffee ☕

Buy me a coffee

Was this page helpful?

Let us know how we did