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.
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.
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.
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.
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.
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_weightsto replace all weights in a matching layer. - Use variable
assignwhen 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
- TensorFlow How can I evaluate a validation data queue multiple times during training?
- TensorFlow How can I evaluate a validation data queue multiple times during training?
- tensorflow how come gather_nd is differentiable?
- Tensorflow How do I convert a EagerTensor into a numpy array?
- Tensorflow How do you monitor GPU performance during model training in real-time?
- Tensorflow How does tf.get_variable work?
- Tensorflow How to extract attention_scores for graphing?
- Tensorflow How to index a tensor using 2D-index like in numpy
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.