How to freeze model layers in Keras

Transfer-learning projects often need one part of a Keras model to preserve learned features while another part adapts to new labels. The trainable flag controls which layer weights the optimizer is allowed to update during fit().

Set trainable on the layer or nested model before calling compile(). A model-level setting propagates to child layers, and compile() records the trainable variables used by the training function, so changing trainable later requires another compile call.

A small nested model is enough to prove the freeze boundary. Freeze the feature base, compile the combined model, train a new head, and compare the base and head weights before accepting the training setup.

1. Create freeze_layers.py with backend selection, imports, and sample arrays.

   freeze_layers.py

   import os
    
   os.environ["KERAS_BACKEND"] = "jax"
    
   import keras
   import numpy as np
    
   keras.utils.set_random_seed(17)
    
   x_train = np.array(
       [
           [0.0, 0.0],
           [0.0, 1.0],
           [1.0, 0.0],
           [1.0, 1.0],
       ],
       dtype="float32",
   )
   y_train = np.array([[0.0], [1.0], [1.0], [0.0]], dtype="float32")

   The KERAS_BACKEND value must be set before import keras. Use the backend already installed for the project.
   Related: How to install Keras with pip
   Related: How to set the Keras backend

2. Add a feature base and a trainable task head.

   base = keras.Sequential(
       [
           keras.layers.Input(shape=(2,)),
           keras.layers.Dense(4, activation="relu", name="feature_dense_1"),
           keras.layers.Dense(3, activation="relu", name="feature_dense_2"),
       ],
       name="feature_base",
   )
    
   inputs = keras.Input(shape=(2,), name="features")
   features = base(inputs)
   outputs = keras.layers.Dense(1, activation="sigmoid", name="task_head")(features)
   model = keras.Model(inputs=inputs, outputs=outputs)

3. Freeze the feature base before compiling the full model.

   base.trainable = False

   Setting trainable on a nested model moves all child-layer weights into the non-trainable set for the compiled training run.

4. Compile the model after setting the frozen state.

   model.compile(
       optimizer=keras.optimizers.SGD(learning_rate=0.5),
       loss=keras.losses.BinaryCrossentropy(),
   )

   Recompile the model after any later trainable change, including unfreezing layers for fine-tuning.
   Related: How to compile a model in Keras

5. Save the base and head weights before training.

   base_before = [weight.copy() for weight in base.get_weights()]
   head_before = [weight.copy() for weight in model.get_layer("task_head").get_weights()]

6. Train the compiled model for a short check run.

   history = model.fit(x_train, y_train, epochs=3, batch_size=4, verbose=0)

7. Compare the base and head weights after training.

   base_after = base.get_weights()
   head_after = model.get_layer("task_head").get_weights()
    
   base_changed = any(
       not np.allclose(before, after) for before, after in zip(base_before, base_after)
   )
   head_changed = any(
       not np.allclose(before, after) for before, after in zip(head_before, head_after)
   )
    
   print(f"backend: {keras.config.backend()}")
   print(f"base.trainable: {base.trainable}")
   print(f"base child flags: {[layer.trainable for layer in base.layers]}")
   print(f"trainable variables: {len(model.trainable_variables)}")
   print(f"non-trainable variables: {len(model.non_trainable_variables)}")
   print("epochs completed:", len(history.history["loss"]))
   print(f"base weights changed: {base_changed}")
   print(f"head weights changed: {head_changed}")

8. Run the script and confirm that only the task head changed.

   $ python3 freeze_layers.py
   backend: jax
   base.trainable: False
   base child flags: [False, False]
   trainable variables: 2
   non-trainable variables: 4
   epochs completed: 3
   base weights changed: False
   head weights changed: True