Build Transfer Learning With Python

Our Project

Suppose we have a dataset containing images of cats and dogs. We want our neural network to predict whether a new image contains a cat or a dog.

Our Dataset

images/
│
├── cats/
│   ├── cat1.jpg
│   ├── cat2.jpg
│   └── cat3.jpg
│
└── dogs/
    ├── dog1.jpg
    ├── dog2.jpg
    └── dog3.jpg

The folder names become our class names.

Install TensorFlow

We will use TensorFlow and Keras to build the model.

pip install tensorflow

After installing TensorFlow, we can import it into our Python program.

import tensorflow as tf

Load the Image Dataset

Keras provides a convenient function for loading images from folders.

train_dataset = tf.keras.utils.image_dataset_from_directory(
    "images/",
    image_size=(224, 224),
    batch_size=32
)

Here, "images/" is the directory containing our class folders.

We resize every image to 224 × 224 pixels because this is the input size commonly used with MobileNetV2.

batch_size=32 means the model processes 32 images at a time.

Check the Class Names

We can see which classes were detected from the folders.

print(train_dataset.class_names)

The result might look like:

['cats', 'dogs']

This means our model has two possible output classes.

Load the Pretrained Model

Now we load MobileNetV2 with weights that were already learned from ImageNet.

base_model = tf.keras.applications.MobileNetV2(
    weights="imagenet",
    include_top=False,
    pooling="avg"
)

The important part is weights="imagenet". This tells TensorFlow to use already-trained weights instead of random weights.

include_top=False removes the original ImageNet classification layer because we need our own classifier for cats and dogs.

pooling="avg" converts the extracted feature maps into a compact feature vector.

Freeze the Pretrained Model

At first, we do not want to change the pretrained weights.

base_model.trainable = False

This means MobileNetV2 will act as a feature extractor. Its learned visual knowledge will remain unchanged.

Image
  ↓
MobileNetV2
  ↓
Pretrained Features
  ↓
New Classifier
  ↓
Cat / Dog

Build Our Model

Now we create a new model containing the pretrained model and our own classification layer.

model = tf.keras.Sequential([
    base_model,
    tf.keras.layers.Dense(
        2,
        activation="softmax"
    )
])

The number 2 represents our two classes: cats and dogs.

Softmax converts the output into probabilities.

Example Output

Cat     → 0.91
Dog     → 0.09

The model would therefore predict Cat.

Compile the Model

Before training, we need to configure the model.

model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)

The optimizer controls how the model updates its trainable weights.

The loss function measures how wrong the predictions are. Accuracy tells us how many predictions are correct.

Train the Model

Now we can train the new classifier.

model.fit(
    train_dataset,
    epochs=5
)

Because MobileNetV2 is frozen, the training mainly updates the new Dense classification layer.

Pretrained Layers
       ↓
    FROZEN
       ↓
New Dense Layer
       ↓
    TRAINED

Fine-Tune the Model

If we want the pretrained model to adapt more closely to our dataset, we can fine-tune some of its layers.

base_model.trainable = True

for layer in base_model.layers[:-20]:
    layer.trainable = False

for layer in base_model.layers[-20:]:
    layer.trainable = True

Now the final 20 layers can learn from our cat and dog images.

We should also recompile the model using a very small learning rate.

model.compile(
    optimizer=tf.keras.optimizers.Adam(
        learning_rate=0.00001
    ),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)

Complete Python Code

import tensorflow as tf

# --------------------------------
# 1. Load dataset
# --------------------------------

train_dataset = tf.keras.utils.image_dataset_from_directory(
    "images/",
    image_size=(224, 224),
    batch_size=32
)

print(train_dataset.class_names)


# --------------------------------
# 2. Load pretrained model
# --------------------------------

base_model = tf.keras.applications.MobileNetV2(
    weights="imagenet",
    include_top=False,
    pooling="avg"
)


# --------------------------------
# 3. Freeze pretrained model
# --------------------------------

base_model.trainable = False


# --------------------------------
# 4. Create new model
# --------------------------------

model = tf.keras.Sequential([
    base_model,
    tf.keras.layers.Dense(
        2,
        activation="softmax"
    )
])


# --------------------------------
# 5. Compile model
# --------------------------------

model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)


# --------------------------------
# 6. Train classifier
# --------------------------------

model.fit(
    train_dataset,
    epochs=5
)


# --------------------------------
# 7. Fine-tune
# --------------------------------

base_model.trainable = True

for layer in base_model.layers[:-20]:
    layer.trainable = False

for layer in base_model.layers[-20:]:
    layer.trainable = True


# --------------------------------
# 8. Recompile
# --------------------------------

model.compile(
    optimizer=tf.keras.optimizers.Adam(
        learning_rate=0.00001
    ),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)


# --------------------------------
# 9. Fine-tune model
# --------------------------------

model.fit(
    train_dataset,
    epochs=5
)

Understand the Complete Flow

The entire program can be understood as a sequence of simple steps.

1. Load Images
       ↓
2. Resize Images
       ↓
3. Load MobileNetV2
       ↓
4. Use ImageNet Weights
       ↓
5. Freeze Pretrained Layers
       ↓
6. Add Cat/Dog Classifier
       ↓
7. Compile Model
       ↓
8. Train Classifier
       ↓
9. Unfreeze Last Layers
       ↓
10. Use Small Learning Rate
       ↓
11. Fine-Tune
       ↓
12. Final Model

Why Is This Better Than Training From Scratch?

Training a deep image model from scratch requires a large amount of data, computation, and training time.

TRAIN FROM SCRATCH

Random Weights
      ↓
Learn Edges
      ↓
Learn Shapes
      ↓
Learn Textures
      ↓
Learn Objects
      ↓
Learn Your Classes


TRANSFER LEARNING

Pretrained Weights
      ↓
Already Learned Visual Features
      ↓
Adapt to Your Classes

That is the main practical advantage of Transfer Learning: we start with useful knowledge instead of throwing it away.

Quick Check