Introduction

A common TensorFlow workflow is to reuse only part of an older checkpoint in a newer model. This happens in transfer learning, architecture expansion, or when you want to keep a pretrained backbone but replace the task-specific head.

The important idea is that TensorFlow restores variables by matching objects and variable names in the checkpoint structure. If the new model contains extra variables, or only wants a subset, you usually restore the matching part and then explicitly allow the rest to remain unmatched.

The Safest Pattern: Restore a Named Submodule

The cleanest approach is to checkpoint the reusable submodule separately and restore that same submodule name into the new model.

1import tensorflow as tf
2
3
4class Backbone(tf.keras.Model):
5    def __init__(self):
6        super().__init__()
7        self.dense = tf.keras.layers.Dense(8, activation="relu")
8
9    def call(self, x):
10        return self.dense(x)
11
12
13class OldModel(tf.keras.Model):
14    def __init__(self):
15        super().__init__()
16        self.backbone = Backbone()
17        self.old_head = tf.keras.layers.Dense(1)
18
19    def call(self, x):
20        x = self.backbone(x)
21        return self.old_head(x)
22
23
24old_model = OldModel()
25_ = old_model(tf.random.normal((1, 4)))
26
27ckpt = tf.train.Checkpoint(backbone=old_model.backbone)
28path = ckpt.save("/tmp/backbone_ckpt")

Now restore that backbone into a new model:

1class NewModel(tf.keras.Model):
2    def __init__(self):
3        super().__init__()
4        self.backbone = Backbone()
5        self.new_head = tf.keras.layers.Dense(3)
6
7    def call(self, x):
8        x = self.backbone(x)
9        return self.new_head(x)
10
11
12new_model = NewModel()
13_ = new_model(tf.random.normal((1, 4)))
14
15restore_ckpt = tf.train.Checkpoint(backbone=new_model.backbone)
16status = restore_ckpt.restore(path)
17status.expect_partial()

The backbone variables match and restore, while the new head remains newly initialized.

Why expect_partial() Matters

When the checkpoint and the current object graph do not match exactly, TensorFlow may warn about unmatched variables or unused checkpoint entries. expect_partial() tells TensorFlow that you are intentionally doing a partial restore.

That is an important signal for the transfer-learning case, because partial restoration is not an error there. It is the design.

Subset Restoration with load_weights

If you are using Keras-style saved weights and the layer structure still matches for the reusable part, load_weights can also work.

new_model.load_weights("/tmp/model_weights")

This is simplest when the shared layers have the same topology and ordering. If the structure has changed more substantially, tf.train.Checkpoint is usually more explicit and reliable.

The more the model diverges from the original graph, the more valuable explicit checkpoint object naming becomes.

Build the Variables Before Restoring

A very common issue is trying to restore before the new model’s variables exist. Subclassed models usually create variables lazily on first call.

That is why the examples above run a dummy forward pass first:

_ = new_model(tf.random.normal((1, 4)))

Without that step, TensorFlow may not yet have created the variables you expect to restore into.

Rename Changes Break Matching

Checkpoint restoration depends on matching structure and names. If you renamed the reusable submodule or substantially changed its internal variable layout, the restore will not match automatically.

For example, if the old checkpoint stores backbone/dense/kernel but the new model renamed that branch to encoder/dense/kernel, the names no longer line up.

In those cases, either keep the reusable submodule naming stable or load weights manually at a lower level.

Common Pitfalls

One common mistake is trying to restore a partial model without allowing partial status explicitly. TensorFlow then warns, and developers assume the restore failed completely when only the unmatched new layers were uninitialized.

Another issue is restoring before the model variables exist. Subclassed models need to be built first so the target variables are present.

It is also easy to change layer names or object structure and then expect an old checkpoint to match automatically. Checkpoint restoration is based on names and graph structure, not on vague conceptual similarity.

Finally, do not assume the newly added layers are restored. After a subset restore, those layers still need initialization and training.

Summary

• TensorFlow can restore a subset of variables into a newer model when the matching object names and variables line up.
• 'tf.train.Checkpoint with a named reusable submodule is usually the clearest pattern.'
• 'expect_partial() is important when unmatched variables are intentional.'
• Build subclassed models before restoring so the target variables actually exist.
• Partial restore is ideal for transfer learning, but only the matching part of the model is restored.