Keras
machine learning
model optimization
weight freezing
deep learning

How to dynamically freeze weights after compiling model in Keras?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

In Keras, freezing a layer means setting its trainable attribute to False so its weights are excluded from gradient updates. The critical detail is that if you change trainable after calling compile(), the change does not affect fit() until you compile the model again.

What Freezing Actually Means

Keras separates two related ideas:

  • whether a layer is trainable
  • which variables the compiled training step will update

Changing layer.trainable updates the layer's configuration, but compile() is what captures the trainable variable set for the built training step.

Freeze Before Compile

The simplest pattern is still to freeze layers first and then compile.

python
1import tensorflow as tf
2
3base = tf.keras.applications.MobileNetV2(
4    include_top=False,
5    input_shape=(96, 96, 3),
6    pooling="avg"
7)
8base.trainable = False
9
10model = tf.keras.Sequential([
11    base,
12    tf.keras.layers.Dense(1, activation="sigmoid")
13])
14
15model.compile(
16    optimizer="adam",
17    loss="binary_crossentropy",
18    metrics=["accuracy"]
19)

This is the cleanest setup for the first training phase in transfer learning.

Change trainable, Then Recompile

If you want to freeze or unfreeze layers after an initial training phase, change the flag and compile again before calling fit().

python
1import tensorflow as tf
2
3base = tf.keras.applications.MobileNetV2(
4    include_top=False,
5    input_shape=(96, 96, 3),
6    pooling="avg"
7)
8base.trainable = False
9
10model = tf.keras.Sequential([
11    base,
12    tf.keras.layers.Dense(1, activation="sigmoid")
13])
14
15model.compile(optimizer="adam", loss="binary_crossentropy")
16
17# Phase 1 training happens here
18
19base.trainable = True
20for layer in base.layers[:-20]:
21    layer.trainable = False
22
23model.compile(
24    optimizer=tf.keras.optimizers.Adam(1e-5),
25    loss="binary_crossentropy"
26)

That second compile() is what makes the new trainable set take effect for future training.

Why Recompiling Matters

Without recompiling, it is easy to think a layer has been frozen or unfrozen when the compiled training function is still using the old variable list. This is one of the most common sources of confusion in Keras fine-tuning workflows.

The official Keras transfer learning guidance is explicit about this step: change trainable, then compile() again.

A Practical Fine-Tuning Pattern

Transfer learning usually looks like this:

  1. freeze the pretrained base
  2. train the new top layers
  3. unfreeze some of the base
  4. recompile with a smaller learning rate
  5. continue training carefully

That small learning-rate reduction matters because the newly unfrozen pretrained weights can be damaged quickly by aggressive updates.

What If You Need Truly Dynamic Behavior Mid-Training

If you want to change trainability during a single training loop without going through Keras fit() phases, use a custom training loop with GradientTape. Then you can choose exactly which variables get gradients on each step.

python
1import tensorflow as tf
2
3layer1 = tf.keras.layers.Dense(4)
4layer2 = tf.keras.layers.Dense(1)
5optimizer = tf.keras.optimizers.Adam()
6
7x = tf.random.normal((8, 3))
8y = tf.random.normal((8, 1))
9
10with tf.GradientTape() as tape:
11    h = layer1(x)
12    preds = layer2(h)
13    loss = tf.reduce_mean(tf.square(preds - y))
14
15trainable_vars = layer2.trainable_variables
16grads = tape.gradient(loss, trainable_vars)
17optimizer.apply_gradients(zip(grads, trainable_vars))

Here only layer2 is updated, regardless of what else exists in the model.

Common Pitfalls

The biggest pitfall is setting layer.trainable = False after compile() and assuming the next fit() will respect it automatically. It will not until the model is compiled again.

Another issue is unfreezing a large pretrained block and continuing with the same learning rate used for training the top layers. That often destabilizes the model.

Developers also forget that batch normalization layers have special behavior in transfer learning. Their trainability and inference/training mode behavior deserve extra care.

Finally, do not overcomplicate the workflow. Most fine-tuning problems are easier to solve with clear phase boundaries than with highly dynamic freezing logic.

Summary

  • Set layer.trainable to freeze or unfreeze weights in Keras.
  • If you change trainability after compile(), call compile() again before fit().
  • A phased transfer-learning workflow is usually simpler than mid-training dynamic changes.
  • Use a smaller learning rate when fine-tuning unfrozen pretrained layers.
  • Use a custom training loop only when you need per-step control over which variables update.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.