TensorFlow Transfer Learning Tutorial for Beginners

By /

Unlocking Deep Learning Power: A Beginner’s Guide to TensorFlow Transfer Learning

Deep Learning, particularly Convolutional Neural Networks (CNNs), has revolutionized how computers “see” and interpret images. From identifying objects in photos to powering self-driving cars and medical diagnoses, image classification and related tasks are central to many modern AI applications. However, training these powerful models from scratch often requires two things that beginners (and even seasoned practitioners) might lack:

  1. Massive Datasets: State-of-the-art models are typically trained on datasets like ImageNet, which contains over 14 million labeled images across thousands of categories. Acquiring and labeling such vast datasets for specific tasks is often infeasible.
  2. Significant Computational Resources: Training deep networks demands substantial processing power (often high-end GPUs or TPUs) and considerable time, potentially days or even weeks.

This is where Transfer Learning comes to the rescue. It’s a powerful technique in machine learning that allows you to leverage knowledge gained from solving one problem and apply it to a different but related problem. In the context of computer vision, it means taking a model already trained on a large dataset (like ImageNet) and adapting it for your specific task, even if you have significantly less data.

This article provides a comprehensive, step-by-step tutorial for beginners on how to perform Transfer Learning for image classification using TensorFlow, Google’s popular open-source machine learning framework, and its user-friendly Keras API. We will cover the core concepts, walk through a practical example (classifying images of cats and dogs), and explore different strategies like feature extraction and fine-tuning.

Target Audience: This tutorial is designed for individuals with basic Python programming knowledge who are new to deep learning or TensorFlow and want to learn how to build effective image classifiers without needing massive datasets or computational resources. Some familiarity with the basic concepts of neural networks is helpful but not strictly required, as we will explain the necessary ideas along the way.

What You Will Learn:

  • What Transfer Learning is and why it’s so effective.
  • The concept of pre-trained models and where they come from.
  • How to choose and load a pre-trained model in TensorFlow/Keras.
  • The difference between Feature Extraction and Fine-Tuning strategies.
  • Step-by-step implementation:
    • Setting up the environment.
    • Preparing and loading your custom dataset (Cats vs. Dogs).
    • Performing Data Augmentation.
    • Building a new model on top of a pre-trained base.
    • Training the model using Feature Extraction.
    • Evaluating model performance.
    • (Optional but recommended) Fine-tuning the model for potentially better results.
    • Making predictions on new images.

By the end of this tutorial, you’ll have a solid understanding of Transfer Learning principles and the practical skills to apply it to your own image classification projects using TensorFlow.

Table of Contents

  1. Understanding Transfer Learning
    • The Core Idea: Leveraging Prior Knowledge
    • Why Does It Work for Images? Hierarchical Features
    • Pre-trained Models: Giants’ Shoulders (ImageNet)
    • Popular Pre-trained Architectures (VGG, ResNet, MobileNet, etc.)
    • Transfer Learning Strategies: Feature Extraction vs. Fine-Tuning
  2. Setting Up Your Environment
    • Prerequisites (Python, pip)
    • Installing TensorFlow
    • Recommended Tools (Jupyter Notebook, Google Colab)
    • Checking Your Setup
  3. The Project: Cats vs. Dogs Classification
    • The Dataset
    • Goal
  4. Step-by-Step Implementation
    • Step 0: Imports and Basic Setup
    • Step 1: Data Acquisition and Preparation
      • Downloading the Dataset
      • Understanding the Directory Structure
      • Using image_data_from_directory
      • Data Preprocessing and Normalization
      • Data Augmentation
      • Creating Data Generators (Datasets)
      • Visualizing the Data
    • Step 2: Choosing and Loading the Pre-trained Base Model
      • Selecting a Base Model (MobileNetV2)
      • Instantiating the Base Model (Without Top Layer)
      • Freezing the Base Model Weights
      • Inspecting the Base Model
    • Step 3: Building the Custom Classifier Head
      • Adding Global Average Pooling
      • Adding the Final Dense Layer(s)
      • Creating the Final Model
      • Inspecting the New Model
    • Step 4: Compiling the Model
      • Choosing an Optimizer
      • Choosing a Loss Function
      • Choosing Metrics
    • Step 5: Training the Model (Feature Extraction Phase)
      • Using model.fit()
      • Monitoring Training Progress
    • Step 6: Evaluating the Feature Extraction Model
      • Plotting Accuracy and Loss Curves
      • Interpreting the Curves (Overfitting/Underfitting)
      • Evaluating on the Validation Set
    • Step 7: Fine-Tuning (Optional but Powerful)
      • The Concept of Fine-Tuning
      • Unfreezing Layers in the Base Model
      • Re-compiling the Model (Lower Learning Rate!)
      • Continuing Training (Fine-Tuning Phase)
      • Evaluating the Fine-Tuned Model
    • Step 8: Making Predictions on New Images
      • Loading and Preprocessing a Single Image
      • Getting Model Predictions
      • Interpreting the Output
  5. Beyond the Basics
    • Trying Other Pre-trained Models
    • TensorFlow Hub
    • Handling More Complex Datasets
    • Other Computer Vision Tasks
  6. Conclusion and Next Steps

1. Understanding Transfer Learning

The Core Idea: Leveraging Prior Knowledge

Imagine you want to learn to identify different types of bicycles (mountain bikes, road bikes, cruisers). You could start from absolute zero, trying to figure out what constitutes a “wheel,” “handlebar,” or “pedal.” This would take a very long time.

Alternatively, you already possess a vast amount of visual knowledge about the world. You know what edges, shapes, textures, and common objects look like. You can leverage this existing knowledge. You already recognize wheels and handlebars from other contexts (cars, motorcycles, scooters). Now, you only need to learn the specific combinations and variations of these features that define different bicycle types. This is much faster and more efficient.

Transfer Learning in deep learning works on a similar principle. Instead of training a neural network from scratch (with randomly initialized weights), we start with a network whose weights have already been trained on a large, general dataset (like ImageNet). This pre-trained network has already learned to recognize a rich hierarchy of visual features – edges, corners, textures, patterns, object parts, and even whole objects. We then adapt this pre-learned knowledge to our specific, often smaller, target dataset.

Why Does It Work for Images? Hierarchical Features

Convolutional Neural Networks (CNNs), the standard architecture for image tasks, learn features hierarchically:

  • Early Layers: Learn simple, generic features like edges, corners, and color blobs. These are universally useful for almost any visual task.
  • Middle Layers: Combine these simple features to learn more complex patterns, textures, and object parts (like eyes, noses, wheels, text). These are still quite general.
  • Later Layers: Combine these parts to recognize specific objects (like dogs, cats, cars, faces). These features become progressively more specialized to the dataset the model was originally trained on.

Transfer learning exploits this hierarchy. We assume that the features learned by the early and middle layers of a model trained on a large dataset (like ImageNet) are general enough to be useful for our specific task (e.g., classifying cats vs. dogs). We keep these layers and their learned weights (often “freezing” them so they don’t change during initial training) and only train the final layers (or add new ones) to learn the specific combinations relevant to our new dataset.

Pre-trained Models: Giants’ Shoulders (ImageNet)

The foundation of transfer learning in computer vision is the availability of models pre-trained on large benchmark datasets. The most famous is ImageNet (specifically the ILSVRC challenge dataset), which contains over 1.2 million training images, 50,000 validation images, and 100,000 test images, categorized into 1000 distinct classes (like different breeds of dogs, types of vehicles, household objects, etc.).

Training models on ImageNet is computationally expensive, but thankfully, researchers and organizations often release the trained weights of their state-of-the-art models. By using these pre-trained models, we are essentially standing on the shoulders of giants, benefiting from the enormous computational effort and vast data used to train them.

Popular Pre-trained Architectures

TensorFlow’s Keras API provides easy access to many popular pre-trained models:

  1. VGG (VGG16, VGG19): Developed by the Visual Geometry Group at Oxford. Known for its simple and uniform architecture (stacks of 3×3 convolution layers). Good performance but quite large (many parameters).
  2. ResNet (ResNet50, ResNet101, ResNet152, ResNet50V2, etc.): Introduced Residual Connections (“shortcuts”) that allow training much deeper networks effectively, mitigating the vanishing gradient problem. Often provides a strong baseline performance.
  3. Inception (InceptionV3, InceptionResNetV2): Developed at Google. Uses “Inception modules” which perform convolutions at different scales in parallel within the same layer, capturing features more efficiently. Known for good performance with relatively fewer parameters than VGG.
  4. MobileNet (MobileNet, MobileNetV2, MobileNetV3): Also from Google, designed specifically for mobile and resource-constrained environments. Uses depthwise separable convolutions to drastically reduce the number of parameters and computational cost while maintaining reasonable accuracy. Excellent choice for applications where speed and model size are critical.
  5. EfficientNet (EfficientNetB0-B7): A family of models that achieves state-of-the-art accuracy by systematically scaling network depth, width, and resolution using a compound scaling method. Often provides the best accuracy-per-parameter trade-off.

The choice of model depends on your specific needs:

  • Need highest accuracy? Consider EfficientNet, InceptionResNetV2, or deeper ResNets.
  • Need speed and small size? MobileNet is a great choice.
  • Good general baseline? ResNet50V2 or InceptionV3 are often solid starting points.

For this tutorial, we will use MobileNetV2 because it offers a good balance of performance and efficiency, making it faster to train, especially without a high-end GPU.

Transfer Learning Strategies: Feature Extraction vs. Fine-Tuning

There are two main ways to apply transfer learning:

  1. Feature Extraction:

    • What: Take the pre-trained model’s convolutional base (all layers except the final fully connected classifier layers) and use it as a fixed feature extractor. Run your new data through it, and train only a new classifier (usually one or two Dense layers) on top of the output features.
    • How: Instantiate the base model with include_top=False. Freeze the weights of the base model (base_model.trainable = False). Add your new classifier layers. Compile and train the model.
    • When: Good starting point, especially when your new dataset is small or very different from the original dataset (though similarity helps). Computationally cheaper as most weights are frozen.

  2. Fine-Tuning:

    • What: Start with Feature Extraction. Then, “unfreeze” some of the top layers of the pre-trained base model and train them along with the newly added classifier layers on your data. The idea is to slightly adjust the pre-trained features (especially the more specialized later ones) to be more relevant to your specific task.
    • How: First, perform feature extraction. Then, set base_model.trainable = True. Optionally, choose a point in the base model (fine_tune_at) and freeze all layers below it. Re-compile the model, crucially using a very low learning rate. Continue training.
    • When: Used when you want to squeeze out potentially better performance, especially if your dataset is reasonably large and somewhat similar to the original dataset (e.g., ImageNet). Requires careful handling (low learning rate!) to avoid destroying the pre-trained weights.

We will implement both strategies in this tutorial.

2. Setting Up Your Environment

Prerequisites

  • Python: You need Python installed (version 3.7-3.11 recommended for current TensorFlow versions). You can download it from python.org.
  • pip: Python’s package installer, usually included with Python installations.

Installing TensorFlow

Open your terminal or command prompt and install TensorFlow using pip:

bash
pip install tensorflow

This command installs the latest stable CPU version of TensorFlow. If you have a compatible NVIDIA GPU and want to leverage it for faster training (highly recommended for deep learning), you’ll need to install the GPU version and necessary drivers/libraries (CUDA, cuDNN). Follow the detailed instructions on the official TensorFlow website: Install TensorFlow with pip.

We also need Matplotlib for plotting graphs:

bash
pip install matplotlib

Recommended Tools

  • Jupyter Notebook / JupyterLab: Interactive environments perfect for experimenting with code, visualizing data, and documenting your workflow step-by-step. Install using pip:
    bash
    pip install notebook # For classic Notebook
    pip install jupyterlab # For JupyterLab
    Then run jupyter notebook or jupyter lab in your terminal.
  • Google Colaboratory (Colab): A free, cloud-based Jupyter Notebook environment provided by Google. It requires no setup and provides free access to GPUs and TPUs, making it ideal for beginners or those without powerful local hardware. Visit colab.research.google.com.

This tutorial assumes you are using an environment like Jupyter or Colab.

Checking Your Setup

Create a new notebook or Python script and run the following code to verify your TensorFlow installation and check for GPU availability:

“`python
import tensorflow as tf
import sys

print(f”Python Version: {sys.version}”)
print(f”TensorFlow Version: {tf.version}”)

Check for GPU

gpu_devices = tf.config.list_physical_devices(‘GPU’)
if gpu_devices:
print(f”Num GPUs Available: {len(gpu_devices)}”)
# Optional: Print details of each GPU
for gpu in gpu_devices:
print(f” – {gpu}”)
else:
print(“No GPU detected. TensorFlow will run on CPU.”)
print(“Note: Training deep learning models on CPU can be very slow.”)
“`

If you have a GPU configured correctly, you should see details about it. Otherwise, it will indicate that it’s using the CPU.

3. The Project: Cats vs. Dogs Classification

The Dataset

We’ll use the popular “Cats vs. Dogs” dataset, originally made available by Microsoft Research for a Kaggle competition. It contains thousands of images of cats and dogs. This is a classic binary image classification problem, perfect for learning transfer learning concepts.

The dataset is typically structured into train and validation (or test) sets, with subdirectories for each class (cat and dog).

Goal

Our objective is to build a TensorFlow model that can accurately classify new images as containing either a cat or a dog, leveraging a pre-trained model (MobileNetV2) via transfer learning. We aim for good validation accuracy, demonstrating the power of transfer learning even with a moderately sized dataset compared to ImageNet.

4. Step-by-Step Implementation

Let’s dive into the code. Follow these steps in your Jupyter Notebook or Colab environment.

Step 0: Imports and Basic Setup

First, import the necessary libraries.

“`python
import matplotlib.pyplot as plt
import numpy as np
import os
import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential

print(f”Using TensorFlow version: {tf.version}”)
“`

Step 1: Data Acquisition and Preparation

Downloading the Dataset

TensorFlow Datasets (tfds) provides easy access to many datasets, but for pedagogical purposes and to show manual handling, we’ll download and extract the Cats vs. Dogs dataset from a URL. TensorFlow provides a utility function get_file for this.

“`python

Define the URL for the dataset

_URL = ‘https://storage.googleapis.com/mledu-datasets/cats_and_dogs_filtered.zip’

Download and extract the dataset

path_to_zip will be the path to the downloaded zip file

path_to_archive will be the path to the directory where it’s extracted

try:
path_to_zip = tf.keras.utils.get_file(
‘cats_and_dogs.zip’, # Name to save the file as
origin=_URL, # URL to download from
extract=True # Automatically extract the contents
)
print(f”Dataset downloaded and extracted to: {path_to_zip}”)
# The extracted directory is usually in the same location as the zip, without the .zip extension
PATH = os.path.join(os.path.dirname(path_to_zip), ‘cats_and_dogs_filtered’)
print(f”Base dataset path: {PATH}”)
except Exception as e:
print(f”Error downloading or extracting dataset: {e}”)
# Handle the error appropriately, maybe exit or try an alternative path
PATH = None # Set PATH to None if download failed

Define paths to the training and validation directories if download was successful

if PATH:
train_dir = os.path.join(PATH, ‘train’)
validation_dir = os.path.join(PATH, ‘validation’)
print(f”Training directory: {train_dir}”)
print(f”Validation directory: {validation_dir}”)
else:
print(“Cannot proceed without the dataset.”)

You can optionally list the contents to verify

!ls {PATH}

!ls {train_dir}

!ls {validation_dir}

“`

Understanding the Directory Structure

The get_file function downloads and extracts the data. The cats_and_dogs_filtered directory should now contain:

cats_and_dogs_filtered/
├── train/
│ ├── cats/ (...)
│ └── dogs/ (...)
└── validation/
├── cats/ (...)
└── dogs/ (...)

This structure is crucial because Keras utilities like image_data_from_directory expect it. Each subdirectory within train and validation corresponds to a class label.

Let’s count the number of images:

“`python
if PATH:
num_cats_tr = len(os.listdir(os.path.join(train_dir, ‘cats’)))
num_dogs_tr = len(os.listdir(os.path.join(train_dir, ‘dogs’)))
num_cats_val = len(os.listdir(os.path.join(validation_dir, ‘cats’)))
num_dogs_val = len(os.listdir(os.path.join(validation_dir, ‘dogs’)))

total_train = num_cats_tr + num_dogs_tr
total_val = num_cats_val + num_dogs_val

print(f'Total training cat images: {num_cats_tr}')
print(f'Total training dog images: {num_dogs_tr}')
print('--')
print(f'Total validation cat images: {num_cats_val}')
print(f'Total validation dog images: {num_dogs_val}')
print('--')
print(f"Total training images: {total_train}")
print(f"Total validation images: {total_val}")

else:
print(“Dataset path not available.”)
“`

You should see 1000 training images and 500 validation images for each class, totaling 2000 training and 1000 validation images. This is a relatively small dataset, making transfer learning highly beneficial.

Using image_data_from_directory

This Keras utility is perfect for loading images from disk when they are organized in the structure we have. It automatically infers class labels from the directory names and generates batches of image data and labels.

First, let’s define some parameters:

python
BATCH_SIZE = 32 # Number of images to process in each batch
IMG_HEIGHT = 160 # Target height for resizing images
IMG_WIDTH = 160 # Target width for resizing images
IMG_SHAPE = (IMG_HEIGHT, IMG_WIDTH, 3) # Expected input shape for MobileNetV2 (3 color channels)

Note: We choose 160×160 as the image size because MobileNetV2 (like many ImageNet models) performs well with specific input sizes (often 224×224, but 160×160 and 192×192 are also common, especially for MobileNet variants). Check the documentation for your chosen base model.

Now, create the data generators. We will not apply data augmentation yet, but we will apply the necessary preprocessing for MobileNetV2 later.

“`python

Create training dataset generator

We will handle preprocessing within the model definition later

train_dataset = tf.keras.utils.image_data_from_directory(
train_dir,
shuffle=True, # Shuffle the data for better training
batch_size=BATCH_SIZE,
image_size=(IMG_HEIGHT, IMG_WIDTH),
label_mode=’binary’ # Since we have only two classes (cats, dogs)
)

Create validation dataset generator

validation_dataset = tf.keras.utils.image_data_from_directory(
validation_dir,
shuffle=False, # No need to shuffle validation data
batch_size=BATCH_SIZE,
image_size=(IMG_HEIGHT, IMG_WIDTH),
label_mode=’binary’
)
“`

image_data_from_directory returns a tf.data.Dataset object. These are highly efficient pipelines for handling data in TensorFlow.

Let’s check the class names inferred by the generator:

python
class_names = train_dataset.class_names
print(f"Class names found: {class_names}") # Should be ['cats', 'dogs']

Data Preprocessing and Normalization

Pre-trained models expect input images to be preprocessed in a specific way, typically matching how the images were processed during the original ImageNet training. For MobileNetV2 (and many others), this involves scaling pixel values from the [0, 255] range to [-1, 1].

Instead of doing this in the data generator, it’s often better practice to include the preprocessing step inside the model itself using Keras preprocessing layers or the model’s dedicated preprocess_input function. This ensures that when you deploy or share your model, the preprocessing is inherently part of it, reducing potential errors.

MobileNetV2 has a specific preprocessing function tf.keras.applications.mobilenet_v2.preprocess_input. We will incorporate this later when building the model.

Data Augmentation

Data augmentation artificially increases the diversity of your training data by applying random transformations (like rotation, flipping, zooming) to the existing images. This helps the model generalize better and reduces overfitting, especially crucial when working with smaller datasets.

Keras provides preprocessing layers for data augmentation that can be added directly into your model definition. This is efficient as the augmentation happens on the GPU during training.

“`python

Define data augmentation layers

data_augmentation = tf.keras.Sequential([
layers.RandomFlip(“horizontal”),
layers.RandomRotation(0.2),
# You could add more augmentations like RandomZoom, RandomContrast etc.
])
“`

We will incorporate this data_augmentation layer into our model after the input layer but before the base model. It should only be active during training, not during inference/validation. Keras layers handle this automatically.

Creating Data Generators (Datasets) – Buffering and Prefetching

For performance, it’s good practice to configure the tf.data.Dataset objects to use buffered prefetching. This allows the data loading to happen asynchronously on the CPU while the GPU is busy training, preventing bottlenecks.

“`python
AUTOTUNE = tf.data.AUTOTUNE

train_dataset = train_dataset.prefetch(buffer_size=AUTOTUNE)
validation_dataset = validation_dataset.prefetch(buffer_size=AUTOTUNE)

If you had a separate test dataset, you’d prefetch it too.

test_dataset = test_dataset.prefetch(buffer_size=AUTOTUNE)

“`

Visualizing the Data

Let’s look at a few images from a training batch to see what they look like after loading (before augmentation/preprocessing within the model).

“`python
plt.figure(figsize=(10, 10))

Take one batch from the training dataset

for images, labels in train_dataset.take(1):
for i in range(9): # Display first 9 images in the batch
ax = plt.subplot(3, 3, i + 1)
# Images loaded by image_data_from_directory are float32 in [0, 255]
# We need to cast to uint8 for display
plt.imshow(images[i].numpy().astype(“uint8″))
plt.title(f”Class: {class_names[int(labels[i])]}”)
plt.axis(“off”)
plt.suptitle(“Sample Training Images”)
plt.show()
“`

You should see a grid of 9 images, correctly labeled as ‘cats’ or ‘dogs’.

Step 2: Choosing and Loading the Pre-trained Base Model

Now, we select and load our pre-trained model. We’ll use MobileNetV2 trained on ImageNet.

Selecting a Base Model (MobileNetV2)

As discussed, MobileNetV2 provides a good trade-off between size, speed, and accuracy.

Instantiating the Base Model (Without Top Layer)

We need to instantiate the MobileNetV2 model from tf.keras.applications. Crucially:

  • input_shape: Must match the shape of our preprocessed images (160, 160, 3).
  • include_top=False: This is vital! It means we only load the convolutional base, excluding the final 1000-neuron Dense layer used for ImageNet classification. We will add our own classifier.
  • weights='imagenet': Specifies that we want to load the weights pre-trained on ImageNet. If set to None, weights would be randomly initialized (defeating the purpose of transfer learning).

“`python

Create the base model from the pre-trained MobileNetV2

base_model = tf.keras.applications.MobileNetV2(
input_shape=IMG_SHAPE,
include_top=False, # <<< Important! Exclude the ImageNet classifier head
weights=’imagenet’
)
“`

TensorFlow will download the MobileNetV2 weights if you haven’t used them before.

Freezing the Base Model Weights

For the initial Feature Extraction phase, we don’t want to update the weights of the pre-trained MobileNetV2 layers during training. We want to preserve the features learned on ImageNet.

“`python

Freeze the convolutional base

base_model.trainable = False
print(f”Base model trainable status set to: {base_model.trainable}”)
“`

Inspecting the Base Model

Let’s look at the architecture of the base model.

“`python

Print a summary of the base model

print(“\n— Base Model Summary —“)
base_model.summary()
“`

You’ll see a long list of layers (convolutional, batch normalization, ReLU activations, etc.) that make up MobileNetV2. Notice the large number of parameters, but importantly, when we check the full model later, these will be listed as “non-trainable parameters” because we froze the base. The output shape of the base model (before the global average pooling we’ll add) will likely be something like (None, 5, 5, 1280) where None represents the batch size, 5x5 is the spatial dimension of the feature map, and 1280 is the number of features extracted.

Step 3: Building the Custom Classifier Head

Now we build our own classifier layers to put on top of the frozen base model. This classifier will take the features extracted by base_model and learn to map them to our two classes (cats, dogs).

Adding Global Average Pooling

The output of base_model is a 3D feature map (5x5x1280 in this case). Before feeding this into a standard Dense (fully connected) layer, we need to flatten it into a 1D vector. A common and effective way to do this is using GlobalAveragePooling2D. It averages the values across the spatial dimensions (5x5), resulting in a single vector per image (1280 elements). This significantly reduces the number of parameters compared to using a Flatten layer followed by a large Dense layer, helping to prevent overfitting.

python
global_average_layer = tf.keras.layers.GlobalAveragePooling2D()

Adding the Final Dense Layer(s)

This is the core of our custom classifier. We need a Dense layer that outputs the final classification score(s).

  • Units: Since this is a binary classification problem (cats vs. dogs), we only need one output unit.
  • Activation: For a single output unit in binary classification, we typically don’t apply an activation function here (the output will be a raw logit). We’ll configure the loss function (BinaryCrossentropy) to expect logits by setting from_logits=True. This is generally more numerically stable than using a sigmoid activation here and setting from_logits=False. If you had more than two classes (multi-class classification), you would use units=num_classes and activation='softmax'.

We can also add a Dropout layer before the final Dense layer for regularization, which randomly sets a fraction of input units to 0 during training to prevent overfitting.

python
dropout_layer = tf.keras.layers.Dropout(0.2) # Dropout 20% of units
prediction_layer = tf.keras.layers.Dense(1) # Single output unit for binary classification

Creating the Final Model

Now, let’s assemble the full model using the Keras Functional API (which is more flexible than Sequential when dealing with pre-trained bases). We need to explicitly define the input, route it through the layers, and define the output.

Remember to include:
1. The input layer.
2. The data augmentation layer (active only during training).
3. The MobileNetV2 preprocessing layer.
4. The frozen base_model.
5. The GlobalAveragePooling2D layer.
6. The optional Dropout layer.
7. The final prediction_layer.

“`python

Define the input layer

inputs = tf.keras.Input(shape=IMG_SHAPE, name=”input_layer”)

Apply data augmentation (active only during training)

x = data_augmentation(inputs)

Apply MobileNetV2 preprocessing

Scale pixel values from [0, 255] to [-1, 1]

preprocess_input = tf.keras.applications.mobilenet_v2.preprocess_input
x = preprocess_input(x)

Pass input through the frozen base model

training=False ensures Batch Normalization layers run in inference mode

(using moving averages) even when the outer model is training.

This is crucial when the base model is frozen.

x = base_model(x, training=False)

Apply global average pooling

x = global_average_layer(x)

Apply dropout for regularization

x = dropout_layer(x) # Active only during training

Apply the final prediction layer

outputs = prediction_layer(x)

Create the final model

model = tf.keras.Model(inputs, outputs)

print(“\n— Full Model Assembled (Feature Extraction) —“)
“`

Inspecting the New Model

Let’s look at the summary of our complete model. Pay close attention to the number of trainable vs. non-trainable parameters.

python
model.summary()

You should see:
* The layers we added (Input, Sequential for augmentation, Lambda for preprocessing, the MobileNetV2 base, GlobalAveragePooling2D, Dropout, Dense).
* A large number of Non-trainable params corresponding to the frozen MobileNetV2 weights.
* A relatively small number of Trainable params corresponding only to the weights in the GlobalAveragePooling2D (none), Dropout (none), and the final Dense layer (plus biases). These are the only weights that will be updated during the feature extraction phase.

Step 4: Compiling the Model

Before we can train the model, we need to configure the learning process using model.compile(). This involves specifying:

  • Optimizer: The algorithm used to update the model weights based on the gradients. Adam is a popular and generally effective choice. We’ll start with a typical learning rate.
  • Loss Function: Measures how inaccurate the model is during training. Since we have binary classification and our output layer produces logits (no activation), we use BinaryCrossentropy(from_logits=True). If we had used activation='sigmoid' in the last layer, we would set from_logits=False. For multi-class, use CategoricalCrossentropy (if labels are one-hot encoded) or SparseCategoricalCrossentropy (if labels are integers).
  • Metrics: Used to monitor the training and testing steps. For classification, accuracy is the most common metric.

“`python
base_learning_rate = 0.001 # A common starting learning rate

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=base_learning_rate),
loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
metrics=[‘accuracy’])

print(“\n— Model Compiled —“)
print(f”Optimizer: Adam, Learning Rate: {base_learning_rate}”)
print(f”Loss Function: Binary Crossentropy (from_logits=True)”)
print(f”Metrics: Accuracy”)
“`

Step 5: Training the Model (Feature Extraction Phase)

Now we are ready to train the model using the model.fit() method. We provide the training data, the number of epochs (passes through the entire dataset), and the validation data to monitor performance on unseen data after each epoch.

“`python
initial_epochs = 10 # Start with a reasonable number of epochs

print(f”\n— Starting Training (Feature Extraction) for {initial_epochs} Epochs —“)

Train the model

history = model.fit(train_dataset,
epochs=initial_epochs,
validation_data=validation_dataset)

print(“\n— Training Finished —“)
“`

Watch the output during training. You’ll see the loss and accuracy on the training set (loss, accuracy) and the validation set (val_loss, val_accuracy) reported after each epoch. Ideally, accuracy should increase and loss should decrease on both sets. Since we froze the vast majority of the network, this phase should train relatively quickly.

Step 6: Evaluating the Feature Extraction Model

Let’s visualize the training and validation performance over the epochs. The history object returned by model.fit() contains this information.

Plotting Accuracy and Loss Curves

“`python

Get accuracy and loss values from history

acc = history.history[‘accuracy’]
val_acc = history.history[‘val_accuracy’]
loss = history.history[‘loss’]
val_loss = history.history[‘val_loss’]
epochs_range = range(initial_epochs)

Plot Training and Validation Accuracy

plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label=’Training Accuracy’)
plt.plot(epochs_range, val_acc, label=’Validation Accuracy’)
plt.legend(loc=’lower right’)
plt.title(‘Training and Validation Accuracy’)
plt.xlabel(‘Epoch’)
plt.ylabel(‘Accuracy’)

Plot Training and Validation Loss

plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label=’Training Loss’)
plt.plot(epochs_range, val_loss, label=’Validation Loss’)
plt.legend(loc=’upper right’)
plt.title(‘Training and Validation Loss’)
plt.xlabel(‘Epoch’)
plt.ylabel(‘Loss’)

plt.suptitle(“Feature Extraction Training Performance”)
plt.show()
“`

Interpreting the Curves (Overfitting/Underfitting)

  • Good Fit: Training and validation accuracy both increase and stabilize at a high value. Training and validation loss both decrease and stabilize at a low value. The gap between training and validation curves is small.
  • Overfitting: Training accuracy keeps increasing (or loss keeps decreasing), while validation accuracy stalls or starts decreasing (or validation loss starts increasing). There’s a significant gap between the training and validation curves. This means the model learned the training data too well, including its noise, and doesn’t generalize to new data. Data augmentation helps mitigate this, but it can still occur.
  • Underfitting: Both training and validation accuracy are low and don’t improve much, or loss remains high. The model is too simple or hasn’t trained long enough to capture the patterns in the data.

Looking at the plots for the feature extraction phase, you’ll likely see good initial improvement. Since the base model is frozen, the model might plateau relatively quickly, and you might observe some overfitting (validation accuracy stalling while training accuracy continues to rise slightly). This suggests that allowing some of the pre-trained weights to adapt to our specific dataset (fine-tuning) could be beneficial.

Evaluating on the Validation Set

We can also get the final loss and accuracy on the validation set using model.evaluate():

python
print("\n--- Evaluating Model on Validation Data (After Feature Extraction) ---")
loss_fe, accuracy_fe = model.evaluate(validation_dataset)
print(f'Validation Loss (Feature Extraction): {loss_fe:.4f}')
print(f'Validation Accuracy (Feature Extraction): {accuracy_fe:.4f}')

You should achieve a decent accuracy (likely >90%, possibly >95%) just with feature extraction, showcasing the power of the pre-trained MobileNetV2 features.

Step 7: Fine-Tuning (Optional but Powerful)

Fine-tuning aims to improve performance further by unfreezing some of the top layers of the base model and training them with a very low learning rate. This allows the model to adapt the more specialized pre-trained features to the nuances of the Cats vs. Dogs dataset.

The Concept of Fine-Tuning

We don’t want to unfreeze the entire base model, especially not the earliest layers that learned very general features (edges, textures). Unfreezing too many layers or using too high a learning rate risks “catastrophic forgetting,” where the valuable pre-trained knowledge is destroyed by large gradient updates driven by our smaller dataset.

We typically unfreeze only the top portion of the base model.

Unfreezing Layers in the Base Model

First, set the entire base model to be trainable:

python
base_model.trainable = True
print(f"\n--- Preparing for Fine-Tuning ---")
print(f"Base model trainable status set to: {base_model.trainable}")

Now, let’s see how many layers are in the base model:

python
print(f"Number of layers in the base model: {len(base_model.layers)}")

We need to decide from which layer onwards we want to fine-tune. Let’s choose to fine-tune from layer 100 onwards (MobileNetV2 has about 154 layers). Layers before this will remain frozen. Note: The optimal fine_tune_at layer index might require experimentation.

“`python

Fine-tune from this layer onwards

fine_tune_at = 100

Freeze all the layers before the fine_tune_at layer

for layer in base_model.layers[:fine_tune_at]:
layer.trainable = False
print(f”Froze layers 0 to {fine_tune_at – 1}”)

Verify which layers are trainable

for i, layer in enumerate(base_model.layers):

print(f”Layer {i}: {layer.name} – Trainable: {layer.trainable}”)

“`

Re-compiling the Model (Lower Learning Rate!)

This is CRITICAL. After changing the trainable status of layers, you MUST re-compile the model for the changes to take effect. Furthermore, when fine-tuning, you should use a much lower learning rate than during feature extraction. This prevents large updates that could wreck the pre-trained weights you’re trying to adapt subtly. A common practice is to use a learning rate about 10x smaller.

“`python

Compile the model with a lower learning rate for fine-tuning

fine_tune_learning_rate = base_learning_rate / 10
model.compile(loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
optimizer=tf.keras.optimizers.RMSprop(learning_rate=fine_tune_learning_rate), # RMSprop is often recommended for fine-tuning
metrics=[‘accuracy’])

print(“\n— Model Re-compiled for Fine-Tuning —“)
print(f”Optimizer: RMSprop, Learning Rate: {fine_tune_learning_rate}”) # Note the much lower LR

Let’s look at the summary again to see trainable parameters now

model.summary()
“`

Now, the model.summary() output should show a significantly larger number of Trainable params, reflecting the unfrzen layers in the MobileNetV2 base, in addition to our classifier head. The number of Non-trainable params will correspond to the layers we explicitly kept frozen (layers 0 to fine_tune_at - 1).

Continuing Training (Fine-Tuning Phase)

We continue training the model, but now both the classifier head and the unfrozen top layers of the base model will be updated. We train for a few more epochs. It’s helpful to use the initial_epoch argument in model.fit to continue the epoch numbering from where the feature extraction phase left off, which makes plotting easier.

“`python
fine_tune_epochs = 10 # Number of epochs for fine-tuning
total_epochs = initial_epochs + fine_tune_epochs

print(f”\n— Starting Fine-Tuning for {fine_tune_epochs} Epochs —“)
print(f”(Continuing from epoch {initial_epochs})”)

Continue training the model

history_fine = model.fit(train_dataset,
epochs=total_epochs,
initial_epoch=history.epoch[-1] + 1, # Start epoch numbering correctly
validation_data=validation_dataset)

print(“\n— Fine-Tuning Finished —“)
“`

Monitor the val_accuracy. You might see it improve further during fine-tuning, potentially surpassing the peak accuracy achieved during feature extraction alone.

Evaluating the Fine-Tuned Model

Let’s update our plots to include the fine-tuning phase.

“`python

Append the fine-tuning history to the initial history

acc += history_fine.history[‘accuracy’]
val_acc += history_fine.history[‘val_accuracy’]
loss += history_fine.history[‘loss’]
val_loss += history_fine.history[‘val_loss’]

Plot combined accuracy curves

plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(range(total_epochs), acc, label=’Training Accuracy’)
plt.plot(range(total_epochs), val_acc, label=’Validation Accuracy’)

Add a vertical line to show where fine-tuning started

plt.plot([initial_epochs-1, initial_epochs-1], plt.ylim(), label=’Start Fine Tuning’)
plt.legend(loc=’lower right’)
plt.title(‘Training and Validation Accuracy (Incl. Fine-Tuning)’)
plt.xlabel(‘Epoch’)
plt.ylabel(‘Accuracy’)

Plot combined loss curves

plt.subplot(1, 2, 2)
plt.plot(range(total_epochs), loss, label=’Training Loss’)
plt.plot(range(total_epochs), val_loss, label=’Validation Loss’)

Add a vertical line

plt.plot([initial_epochs-1, initial_epochs-1], plt.ylim(), label=’Start Fine Tuning’)
plt.legend(loc=’upper right’)
plt.title(‘Training and Validation Loss (Incl. Fine-Tuning)’)
plt.xlabel(‘Epoch’)
plt.ylabel(‘Loss’)

plt.suptitle(“Full Training Performance (Feature Extraction + Fine-Tuning)”)
plt.show()

Evaluate the final fine-tuned model

print(“\n— Evaluating Model on Validation Data (After Fine-Tuning) —“)
loss_ft, accuracy_ft = model.evaluate(validation_dataset)
print(f’Validation Loss (Fine-Tuned): {loss_ft:.4f}’)
print(f’Validation Accuracy (Fine-Tuned): {accuracy_ft:.4f}’)

print(f”\nAccuracy Improvement from Fine-Tuning: {accuracy_ft – accuracy_fe:.4f}”)
“`

You should observe that fine-tuning potentially pushed the validation accuracy slightly higher than feature extraction alone, demonstrating its value for squeezing out extra performance. The plots clearly show the two phases of training.

Step 8: Making Predictions on New Images

Now that we have a trained model, let’s use it to predict the class of a new image it hasn’t seen before.

Loading and Preprocessing a Single Image

We need a function that takes an image file path, loads it, resizes it to the expected IMG_HEIGHTxIMG_WIDTH, converts it to a NumPy array, adds a batch dimension (models expect batches of images), and crucially, applies the same preprocessing used during training (scaling pixels to [-1, 1]).

“`python
def preprocess_image(image_path):
“””Loads and preprocesses an image for model prediction.”””
img = tf.keras.utils.load_img(
image_path, target_size=(IMG_HEIGHT, IMG_WIDTH)
)
img_array = tf.keras.utils.img_to_array(img)
# Add batch dimension: (height, width, channels) -> (1, height, width, channels)
img_array = tf.expand_dims(img_array, 0)

# IMPORTANT: Preprocess the image using the *exact same* function as in the model
# Note: We don't need data augmentation or the initial [0, 255] scaling here,
# as `img_to_array` already gives us floats. The preprocessing is handled by
# the `preprocess_input` function which is part of the model itself.
# However, if preprocessing wasn't part of the model, you'd apply it here.

# Since our model includes the preprocessing layer, we can just feed the
# scaled [0, 255] array directly. The model will handle the rest.
# If your model DID NOT include the preprocess_input layer, you would do:
# preprocessed_img_array = tf.keras.applications.mobilenet_v2.preprocess_input(img_array)
# return preprocessed_img_array
return img_array # Shape: (1, 160, 160, 3), range [0, 255]

Example usage: You need an image file (e.g., ‘my_cat.jpg’ or ‘my_dog.png’)

Make sure the image file exists in your environment.

If using Colab, upload an image file first.

Replace with the actual path to your test image

test_image_path = ‘path/to/your/test_image.jpg’ # <— CHANGE THIS

if os.path.exists(test_image_path):
processed_image = preprocess_image(test_image_path)
print(f”Image loaded and preprocessed. Shape: {processed_image.shape}”)

# Display the test image (optional)
plt.imshow(tf.keras.utils.load_img(test_image_path, target_size=(IMG_HEIGHT, IMG_WIDTH)))
plt.title("Test Image")
plt.axis("off")
plt.show()

else:
print(f”Test image not found at: {test_image_path}”)
processed_image = None

“`

Getting Model Predictions

Use the model.predict() method on the preprocessed image.

“`python
if processed_image is not None:
# Get model predictions (logits)
predictions = model.predict(processed_image)
raw_score = predictions[0][0] # Output is shape (1, 1), get the single logit value

print(f"Raw prediction score (logit): {raw_score:.4f}")

else:
print(“Cannot make prediction, image not loaded.”)

“`

Interpreting the Output

The output raw_score is a logit (because we used BinaryCrossentropy(from_logits=True) and no final activation).

  • Logits > 0 generally correspond to the positive class (which is ‘dogs’ based on the alphabetical order ['cats', 'dogs']).
  • Logits < 0 generally correspond to the negative class (‘cats’).
  • Logits close to 0 indicate uncertainty.

To convert the logit to a probability (between 0 and 1), we can apply the sigmoid function:

“`python
if processed_image is not None:
# Apply sigmoid to convert logit to probability
probability_dog = tf.nn.sigmoid(raw_score).numpy()

print(f"Probability of being a dog: {probability_dog:.4f}")
print(f"Probability of being a cat: {1 - probability_dog:.4f}")

# Determine the predicted class
if probability_dog > 0.5:
    print("Prediction: This is a Dog!")
else:
    print("Prediction: This is a Cat!")

else:
print(“Cannot interpret prediction.”)
“`

Congratulations! You have successfully used transfer learning to train an image classifier and make predictions on new data.

5. Beyond the Basics

This tutorial covered the fundamentals, but there’s much more to explore:

  • Trying Other Pre-trained Models: Experiment with different base models available in tf.keras.applications (ResNet50V2, EfficientNetB0, InceptionV3). Remember to adjust IMG_SHAPE and potentially the preprocessing function according to the model’s requirements. Compare their performance and training time.
  • TensorFlow Hub (tfhub.dev): An online repository of pre-trained models, including many vision models (feature vectors and full models). tensorflow_hub library offers an alternative way to load pre-trained components, often simplifying the process.
  • Handling More Complex Datasets: Apply these techniques to datasets with more classes (multi-class classification). You’ll need to change label_mode in image_data_from_directory to 'categorical', use units=num_classes and activation='softmax' in your final Dense layer, and switch the loss to CategoricalCrossentropy or SparseCategoricalCrossentropy.
  • Different Computer Vision Tasks: Transfer learning isn’t limited to classification. Pre-trained backbones are commonly used as feature extractors for object detection (drawing bounding boxes around objects), semantic segmentation (classifying each pixel in an image), and more.
  • Hyperparameter Tuning: Experiment with different learning rates, optimizer choices, dropout rates, batch sizes, the number of layers to fine-tune (fine_tune_at), and data augmentation strategies to optimize performance. Tools like KerasTuner can help automate this.
  • Deployment: Once you have a trained model, you might want to deploy it. TensorFlow Lite (tf.lite) is designed for deploying models on mobile, embedded, and IoT devices, often requiring model conversion and optimization. TensorFlow Serving is used for deploying models at scale on servers.

6. Conclusion and Next Steps

Transfer Learning is an indispensable technique in the modern deep learning toolkit. It democratizes the power of large-scale models, allowing developers and researchers to build highly effective image classifiers (and other vision models) even with limited data and computational resources.

In this comprehensive tutorial, we’ve walked through:

  1. The concepts behind transfer learning, pre-trained models, feature extraction, and fine-tuning.
  2. Setting up a TensorFlow environment.
  3. Preparing the Cats vs. Dogs dataset using Keras utilities.
  4. Implementing data augmentation.
  5. Loading a pre-trained MobileNetV2 base model and freezing its weights.
  6. Building and compiling a custom classifier on top.
  7. Training the model using the feature extraction strategy.
  8. Evaluating the results and understanding performance graphs.
  9. Implementing the fine-tuning strategy by unfreezing layers and using a lower learning rate to further boost performance.
  10. Making predictions on new images.

You now have a solid foundation and a practical workflow for applying transfer learning to your own image classification problems using TensorFlow.

Next Steps:

  • Practice: Apply this workflow to a different dataset (e.g., tf_flowers, food101 from TensorFlow Datasets, or your own image collection).
  • Experiment: Try different pre-trained models and hyperparameters.
  • Explore: Investigate TensorFlow Hub for more pre-trained model options.
  • Deepen: Learn more about CNN architectures, optimizers, loss functions, and regularization techniques.
  • Expand: Look into applying transfer learning to other tasks like object detection or segmentation.

The world of computer vision is vast and exciting. By mastering techniques like transfer learning, you’ve taken a significant step towards building powerful and practical AI applications. Happy coding!

Leave a Comment

Your email address will not be published. Required fields are marked *

Scroll to Top