How to run transfer learning in Keras

Transfer learning adapts a pretrained Keras application model to a new task by freezing the feature extractor and training a small task-specific head. This saves training time when the new dataset is smaller than the dataset used to train the base model.

The basic workflow is to load a pretrained base with include_top=False, freeze it, add a classifier head, compile, train the head, then optionally unfreeze a small top section for fine-tuning with a lower learning rate.

The sample uses generated image arrays only to validate the wiring and trainability checks. Replace those arrays with a real image dataset before using the workflow to judge model quality.

1. Install Keras and the backend package.

   $ python -m pip install --upgrade keras jax

2. Create run_transfer_learning.py.

   run_transfer_learning.py

   import os
    
   os.environ["KERAS_BACKEND"] = "jax"
    
   import keras
   import numpy as np
   from keras import layers
    
    
   keras.utils.set_random_seed(23)
    
   num_classes = 2
   image_shape = (96, 96, 3)
    
   rng = np.random.default_rng(23)
   x_train = rng.uniform(size=(8, *image_shape)).astype("float32")
   y_train = np.array([0, 1, 0, 1, 0, 1, 0, 1], dtype="int32")
   x_val = rng.uniform(size=(4, *image_shape)).astype("float32")
   y_val = np.array([0, 1, 0, 1], dtype="int32")
    
   base_model = keras.applications.MobileNetV2(
       input_shape=image_shape,
       include_top=False,
       weights="imagenet",
       pooling="avg",
   )
   base_model.trainable = False
    
   inputs = keras.Input(shape=image_shape, name="image")
   x = keras.applications.mobilenet_v2.preprocess_input(inputs)
   x = base_model(x, training=False)
   x = layers.Dropout(0.2)(x)
   outputs = layers.Dense(num_classes, activation="softmax", name="class_probs")(x)
   model = keras.Model(inputs, outputs, name="mobilenet_transfer_demo")
    
   model.compile(
       optimizer=keras.optimizers.Adam(learning_rate=0.001),
       loss=keras.losses.SparseCategoricalCrossentropy(),
       metrics=[keras.metrics.SparseCategoricalAccuracy(name="accuracy")],
   )
   history = model.fit(x_train, y_train, validation_data=(x_val, y_val), epochs=1, batch_size=4, verbose=0)
   frozen_trainable_weights = len(model.trainable_weights)
    
   base_model.trainable = True
   for layer in base_model.layers[:-20]:
       layer.trainable = False
   model.compile(
       optimizer=keras.optimizers.Adam(learning_rate=0.0001),
       loss=keras.losses.SparseCategoricalCrossentropy(),
       metrics=[keras.metrics.SparseCategoricalAccuracy(name="accuracy")],
   )
   fine_tune_history = model.fit(
       x_train,
       y_train,
       validation_data=(x_val, y_val),
       epochs=2,
       initial_epoch=1,
       batch_size=4,
       verbose=0,
   )
   model.save("transfer-learning-demo.keras")
    
   print(f"backend: {keras.backend.backend()}")
   print(f"base model: {base_model.name}")
   print("base weights: imagenet")
   print(f"frozen trainable weights: {frozen_trainable_weights}")
   print(f"fine-tune trainable weights: {len(model.trainable_weights)}")
   print(f"final val accuracy: {fine_tune_history.history['val_accuracy'][-1]:.4f}")
   print("saved model: transfer-learning-demo.keras")

   The first run downloads the pretrained MobileNetV2 weights into the backend's Keras cache. Use a real image dataset and labels in place of the generated arrays.

3. Run the transfer-learning script.

   $ python run_transfer_learning.py
   backend: jax
   base model: mobilenetv2_1.00_96
   base weights: imagenet
   frozen trainable weights: 2
   fine-tune trainable weights: 22
   final val accuracy: 0.5000
   saved model: transfer-learning-demo.keras

   The trainable-weight count should increase only after unfreezing the top base-model layers and recompiling.

4. Save the adapted model after the final fine-tuning phase.

   model.save("transfer-learning-demo.keras")

   Related: How to save and load a Keras model