Transfer Learning & Pretrained Models

# Run this in Google Colab or on your own computer (it downloads the pretrained weights)
import keras
from keras import layers

# 1. Load MobileNetV2, pretrained on over a million ImageNet photos, minus its top layer
base = keras.applications.MobileNetV2(
    input_shape=(160, 160, 3), include_top=False, weights="imagenet")
base.trainable = False                          # freeze its 2.2 million weights

# 2. Load YOUR images: one folder per class, e.g. photos/cracked, photos/ok
train_ds = keras.utils.image_dataset_from_directory(
    "photos", validation_split=0.2, subset="training", seed=1,
    image_size=(160, 160), batch_size=32)
val_ds = keras.utils.image_dataset_from_directory(
    "photos", validation_split=0.2, subset="validation", seed=1,
    image_size=(160, 160), batch_size=32)

# 3. Add a small new head for your classes
inputs = keras.Input(shape=(160, 160, 3))
x = keras.applications.mobilenet_v2.preprocess_input(inputs)   # same scaling as ImageNet
x = base(x, training=False)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dropout(0.2)(x)
outputs = layers.Dense(len(train_ds.class_names), activation="softmax")(x)
model = keras.Model(inputs, outputs)

# 4. Stage 1: train only the new head
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
model.fit(train_ds, validation_data=val_ds, epochs=10)

# 5. Stage 2: fine-tune the top layers of the base with a tiny learning rate
base.trainable = True
for layer in base.layers[:-30]:                 # keep the early layers frozen
    layer.trainable = False
model.compile(optimizer=keras.optimizers.Adam(1e-5),
              loss="sparse_categorical_crossentropy", metrics=["accuracy"])
model.fit(train_ds, validation_data=val_ds, epochs=10)

import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import numpy as np
from sklearn.datasets import load_digits
import keras
from keras import layers

X, y = load_digits(return_X_y=True)
X = (X / 16.0).reshape(-1, 8, 8, 1)

# The "big" task: digits 0-4, with plenty of data (about 900 images)
big = y < 5
X_big, y_big = X[big], y[big]

# The NEW task: digits 5-9, with only 5 labelled images of each
rng = np.random.default_rng(0)
new_idx = np.where(y >= 5)[0]
train_idx = np.concatenate([rng.choice(new_idx[y[new_idx] == d], 5, replace=False)
                            for d in range(5, 10)])
test_idx = np.setdiff1d(new_idx, train_idx)
X_small, y_small = X[train_idx], y[train_idx] - 5
X_test, y_test = X[test_idx], y[test_idx] - 5

def feature_extractor():
    return keras.Sequential([
        keras.Input(shape=(8, 8, 1)),
        layers.Conv2D(32, 3, padding="same", activation="relu"),
        layers.MaxPooling2D(2),
        layers.Conv2D(64, 3, padding="same", activation="relu"),
        layers.Flatten(),
    ], name="features")

def with_new_head(base):
    model = keras.Sequential([base, layers.Dense(5, activation="softmax")])
    model.compile(optimizer=keras.optimizers.Adam(1e-3),
                  loss="sparse_categorical_crossentropy", metrics=["accuracy"])
    return model

# 1. Pretrain on the big task
keras.utils.set_random_seed(0)
base = feature_extractor()
pre = keras.Sequential([base, layers.Dense(5, activation="softmax")])
pre.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
pre.fit(X_big, y_big, epochs=30, verbose=0)
pretrained_weights = base.get_weights()

# 2. Three ways to learn the new task from just 25 images
results = {}
keras.utils.set_random_seed(1)
scratch = with_new_head(feature_extractor())                  # a) start from nothing
scratch.fit(X_small, y_small, epochs=100, verbose=0)
results["From scratch"] = scratch.evaluate(X_test, y_test, verbose=0)[1]

keras.utils.set_random_seed(1)
frozen_base = feature_extractor(); frozen_base.set_weights(pretrained_weights)
frozen_base.trainable = False                                 # b) freeze the base
frozen = with_new_head(frozen_base)
frozen.fit(X_small, y_small, epochs=100, verbose=0)
results["Feature extraction (frozen)"] = frozen.evaluate(X_test, y_test, verbose=0)[1]

frozen_base.trainable = True                                  # c) then fine-tune it all
fine = keras.Sequential([frozen_base, frozen.layers[-1]])
fine.compile(optimizer=keras.optimizers.Adam(1e-4),           # much smaller learning rate
             loss="sparse_categorical_crossentropy", metrics=["accuracy"])
fine.fit(X_small, y_small, epochs=30, verbose=0)
results["Fine-tuned"] = fine.evaluate(X_test, y_test, verbose=0)[1]

print(f"Training images for the new task: {len(X_small)}  |  Test images: {len(X_test)}")
for name, acc in results.items():
    print(f"{name:28} test accuracy = {acc:.3f}")

Training images for the new task: 25  |  Test images: 871
From scratch                 test accuracy = 0.931
Feature extraction (frozen)  test accuracy = 0.892
Fine-tuned                   test accuracy = 0.894