How to get weights from .pb model in Tensorflow
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
A .pb file often stores a frozen TensorFlow graph used for inference, not a training checkpoint. You can still inspect constant tensors inside that graph, including many weight arrays. The extraction workflow depends on whether the file is a frozen GraphDef or a SavedModel export.
Identify the Model Format First
The extension alone is not enough. A standalone frozen graph is usually one .pb file containing nodes and constants. A SavedModel directory contains saved_model.pb plus variables files.
For frozen graphs, weights are typically embedded as Const node values. For SavedModel, weights usually live in variables checkpoints and are easier to read through the loaded model object.
Extract Constants From a Frozen GraphDef
The following script reads a frozen graph and prints tensor names and shapes for constants. It also stores selected values in NumPy format.
This method works when the graph was frozen with variable values converted to constants.
Extract Weights From SavedModel Path
If your .pb is part of a SavedModel directory, use the high level loader and inspect variables directly.
If the model originated from Keras, conversion back to a Keras object may allow get_weights(), but that is not guaranteed for every export path.
Map Tensor Names to Layers
Raw constant names can be cryptic. To make output useful, match names to operation scopes. A common practice is to group by the prefix before the last slash and inspect shape patterns.
This helps you identify kernel tensors, bias vectors, and batch normalization statistics.
Export Extracted Weights for Analysis
After extraction, store tensors in a format your tooling can consume. NumPy binary files are convenient for Python workflows, while comma separated text works for quick inspection in spreadsheets. Keep original tensor names in a manifest file so analysis results can be traced back to graph nodes.
This is useful when comparing two model versions and looking for unexpected parameter drift.
Common Pitfalls
- Assuming every
.pbcontains trainable variables: many inference graphs have only constants. - Mixing TensorFlow major versions: TF2 eager defaults can confuse TF1 style graph loading.
- Expecting layer names from original training code: optimization passes may rename nodes.
- Loading huge tensors into memory at once: extraction scripts can crash on limited machines.
- Ignoring legal and compliance constraints: some model artifacts are licensed for inference only.
Summary
- Determine whether the
.pbis frozen GraphDef or part of SavedModel. - For frozen graphs, read
Constnodes withtensor_util.MakeNdarray. - For SavedModel, inspect loaded variables through TensorFlow APIs.
- Group tensor names by scope to make extracted weights interpretable.
- Keep TensorFlow version compatibility in mind during extraction.
Related reading
- How to get weights from tensorflow fully_connected
- How to get weights in tf.layers.dense?
- How to give a constant input to keras
- How to graph tf.keras model in Tensorflow-2.0?
- how to give the test size in stratified kfold sampling in python?
- How to graph grid scores from GridSearchCV?
- How to handle large amouts of data in tensorflow?
- How to handle large amouts of data in tensorflow?
.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.