TensorFlow
.pb model
weights extraction
machine learning
neural networks

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.

Practice ML system design

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.

python
1import tensorflow as tf
2import numpy as np
3from tensorflow.python.framework import tensor_util
4
5graph_def = tf.compat.v1.GraphDef()
6with tf.io.gfile.GFile("model.pb", "rb") as f:
7    graph_def.ParseFromString(f.read())
8
9constants = {}
10for node in graph_def.node:
11    if node.op == "Const":
12        value = tensor_util.MakeNdarray(node.attr["value"].tensor)
13        constants[node.name] = value
14
15print(f"Found {len(constants)} constant tensors")
16for name, arr in list(constants.items())[:10]:
17    print(name, arr.shape, arr.dtype)
18
19# Example export of one tensor
20if constants:
21    first_name = next(iter(constants))
22    np.save("extracted_weight.npy", constants[first_name])

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.

python
1import tensorflow as tf
2
3loaded = tf.saved_model.load("exported_model")
4
5# If signatures exist, inspect available serving functions
6print(list(loaded.signatures.keys()))
7
8# Many SavedModels expose variables via trackable attributes
9if hasattr(loaded, "variables"):
10    for v in loaded.variables[:10]:
11        print(v.name, v.shape)

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.

python
1from collections import defaultdict
2
3by_scope = defaultdict(list)
4for name, arr in constants.items():
5    scope = name.rsplit("/", 1)[0] if "/" in name else name
6    by_scope[scope].append((name, arr.shape))
7
8for scope, items in list(by_scope.items())[:5]:
9    print("Scope:", scope)
10    for item_name, shape in items[:3]:
11        print("  ", item_name, shape)

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.

python
1import json
2
3manifest = []
4for name, arr in constants.items():
5    file_name = name.replace("/", "_") + ".npy"
6    np.save(file_name, arr)
7    manifest.append({"tensor": name, "file": file_name, "shape": list(arr.shape)})
8
9with open("weights_manifest.json", "w", encoding="utf-8") as f:
10    json.dump(manifest, f, indent=2)

This is useful when comparing two model versions and looking for unexpected parameter drift.

Common Pitfalls

  • Assuming every .pb contains 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 .pb is frozen GraphDef or part of SavedModel.
  • For frozen graphs, read Const nodes with tensor_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
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

All Rights Reserved.