TensorFlow
numpy
machine learning
pre-trained weights
neural networks

Tensorflow How can I assign numpy pre-trained weights to subsections of graph?

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

If you already have pre-trained weights in NumPy arrays, you can load them into only part of a TensorFlow model instead of restoring an entire checkpoint. The stable approach is to target public TensorFlow or Keras objects, build the model first, and then assign weights to the exact layer or variable you want to reuse.

Use set_weights on Matching Layers

For Keras models, the cleanest path is usually to copy weights layer by layer. The source can be another model or a list of NumPy arrays that match the target layer's weight shapes.

python
1import numpy as np
2import tensorflow as tf
3
4source = tf.keras.Sequential([
5    tf.keras.layers.Input(shape=(4,)),
6    tf.keras.layers.Dense(8, activation="relu", name="encoder"),
7    tf.keras.layers.Dense(1, name="head"),
8])
9
10target = tf.keras.Sequential([
11    tf.keras.layers.Input(shape=(4,)),
12    tf.keras.layers.Dense(8, activation="relu", name="encoder"),
13    tf.keras.layers.Dense(3, name="new_head"),
14])
15
16# Build both models.
17sample = tf.zeros((1, 4))
18source(sample)
19target(sample)
20
21# Pretend these came from a saved NumPy source.
22pretrained_weights = source.get_layer("encoder").get_weights()
23
24# Assign only the encoder weights.
25target.get_layer("encoder").set_weights(pretrained_weights)

This transfers only the encoder layer. The new head keeps its own initialization, which is a common transfer-learning pattern.

Assign Individual Variables for Finer Control

Sometimes you do not want to replace all weights in a layer. In that case, assign directly to the underlying variables.

python
1import numpy as np
2import tensorflow as tf
3
4layer = tf.keras.layers.Dense(8, name="encoder")
5layer(tf.zeros((1, 4)))  # Build the layer.
6
7kernel_np = np.ones((4, 8), dtype=np.float32)
8bias_np = np.zeros((8,), dtype=np.float32)
9
10layer.kernel.assign(kernel_np)
11layer.bias.assign(bias_np)

This is useful when you have weight arrays for only one subsection, such as a kernel but not a bias, or when you want to modify only one slice of a larger variable.

Load Weights into Subsections of a Larger Model

If your model is composed of named sublayers, target those sublayers explicitly. That keeps the transfer isolated and readable.

python
1class EncoderClassifier(tf.keras.Model):
2    def __init__(self):
3        super().__init__()
4        self.encoder = tf.keras.layers.Dense(8, activation="relu", name="encoder")
5        self.classifier = tf.keras.layers.Dense(2, name="classifier")
6
7    def call(self, inputs):
8        x = self.encoder(inputs)
9        return self.classifier(x)
10
11
12model = EncoderClassifier()
13model(tf.zeros((1, 4)))  # Build variables.
14
15encoder_kernel = np.full((4, 8), 0.5, dtype=np.float32)
16encoder_bias = np.zeros((8,), dtype=np.float32)
17
18model.encoder.set_weights([encoder_kernel, encoder_bias])

The important part is building the model first. Keras layers do not have concrete weight variables until they have been built by build(...) or by being called with sample input.

Validate Shapes Before Assignment

Shape mismatches are the main failure mode when loading NumPy arrays. Inspect the target variables before assignment and compare them with the arrays you plan to load.

python
1for weight in model.encoder.weights:
2    print(weight.name, weight.shape)
3
4print(encoder_kernel.shape)
5print(encoder_bias.shape)

If a layer was trained with a different input size, hidden size, or bias setting, direct assignment will fail. In that case, you need an architectural match or a deliberate conversion step.

Common Pitfalls

The biggest mistake is trying to load arrays before the target model or layer has been built. set_weights cannot succeed if the receiving layer has not created its variables yet.

Another problem is assuming order without checking it. set_weights expects NumPy arrays in the same order that get_weights returns them. For a standard dense layer, that is usually kernel first, then bias, but you should still verify rather than guess.

Developers also sometimes try to assign weights from an old checkpoint into layers whose shapes no longer match after a model refactor. Matching names are not enough; the array shapes must align too.

Finally, avoid reaching into TensorFlow internal modules for this task. Public Keras methods such as get_weights, set_weights, and variable assign are the stable interfaces for loading NumPy weights.

Summary

  • Build the target model or layer before assigning any NumPy weights.
  • Use set_weights to replace all weights in a matching layer.
  • Use variable assign when you need finer-grained control over a specific tensor.
  • Transfer weights only into the sublayers you want to reuse, such as an encoder.
  • Check array order and shape carefully before loading pre-trained values.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free 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