Model Conversion
.pth to .pb
TensorFlow
PyTorch
Deep Learning

How can we convert a .pth model into .pb file?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

There is no direct one-step ".pth to .pb" exporter in PyTorch. A .pth file usually contains PyTorch weights or a serialized checkpoint, while a .pb file belongs to the TensorFlow ecosystem, so conversion normally requires an intermediate format and careful compatibility checks.

First Clarify What You Actually Have

A .pth file is often just a state_dict, not a complete model definition. That means you cannot convert it by file extension alone. You must reconstruct the original PyTorch model class, load the weights, and switch the model to evaluation mode before any export step.

python
1import torch
2from my_model import MyModel
3
4model = MyModel()
5state = torch.load("model.pth", map_location="cpu")
6model.load_state_dict(state)
7model.eval()

If the checkpoint contains optimizer state or training metadata, you need to extract only the model weights from the saved structure.

Why .pb Is Not the Best Mental Target

In TensorFlow 2, the usual export target is a SavedModel directory, which contains a saved_model.pb file plus variables and metadata. That is different from the older standalone frozen-graph .pb files many tutorials still mention.

So the practical question is usually not "How do I create an arbitrary .pb file," but "How do I move this PyTorch model into a TensorFlow-compatible deployment format?"

Typical Path: PyTorch to ONNX

The modern PyTorch-supported export route is ONNX. Once the model is loaded, export it with an example input tensor.

python
1import torch
2from my_model import MyModel
3
4model = MyModel()
5model.load_state_dict(torch.load("model.pth", map_location="cpu"))
6model.eval()
7
8dummy = torch.randn(1, 3, 224, 224)
9
10torch.onnx.export(
11    model,
12    (dummy,),
13    "model.onnx",
14    input_names=["input"],
15    output_names=["output"],
16    dynamo=True,
17)

This gives you an ONNX graph, which is a much more natural interchange format for PyTorch than TensorFlow protobuf.

Then Convert From ONNX to TensorFlow If You Really Need It

If a TensorFlow runtime is mandatory, the next step is an ONNX-to-TensorFlow conversion tool. This is where compatibility becomes fragile. Not every operator maps cleanly across frameworks, and some community conversion tools have maintenance gaps.

In practice, the workflow is:

  1. export PyTorch to ONNX,
  2. validate the ONNX model,
  3. convert ONNX into a TensorFlow representation,
  4. test outputs against the original PyTorch model.

A command-line conversion may look like this with a community converter:

bash
onnx-tf convert -i model.onnx -o saved_model_dir

That output is typically a TensorFlow SavedModel directory rather than a lone frozen .pb file.

Verify Numerical Equivalence

Do not assume the converted model is correct just because the conversion command succeeds. Run the same sample input through the PyTorch model and the converted artifact, then compare outputs.

python
1import numpy as np
2import onnxruntime as ort
3import torch
4
5torch_output = model(dummy).detach().cpu().numpy()
6
7session = ort.InferenceSession("model.onnx")
8onnx_output = session.run(
9    None,
10    {"input": dummy.numpy()},
11)[0]
12
13print(np.max(np.abs(torch_output - onnx_output)))

Checking the ONNX stage first helps isolate where a mismatch appears.

Sometimes Rebuilding in TensorFlow Is the Better Option

If the model uses custom PyTorch layers, unsupported operators, dynamic control flow, or tracing-unfriendly code, conversion can become more painful than rebuilding the model directly in TensorFlow and loading equivalent weights where feasible.

That is especially true for research models with Python-heavy forward passes. Export tools work best when the computation graph is already close to a framework-neutral tensor program.

Alternatives to .pb

Before committing to conversion, ask what the deployment target actually requires. In many cases:

  • ONNX Runtime can serve the model directly,
  • TensorRT can consume ONNX,
  • or a TorchScript or torch.export artifact may fit the deployment environment better.

If the system only says "we need a portable model file," .pb may be the wrong requirement.

Common Pitfalls

The biggest pitfall is treating .pth as a full model definition. Usually it is only weights, so you must recreate the architecture first.

Another mistake is expecting a guaranteed one-click PyTorch-to-TensorFlow conversion. Operator support, control flow, and custom modules can all break the path.

Developers also often target an old frozen-graph .pb file when what they really need is a TensorFlow SavedModel that includes saved_model.pb.

Finally, never skip validation. Conversion pipelines can succeed syntactically while still producing numerically incorrect outputs.

Summary

  • There is no direct official .pth to .pb export path.
  • Rebuild the PyTorch model, load the .pth weights, and export to ONNX first.
  • If TensorFlow is required, convert the ONNX model into a TensorFlow SavedModel and verify the results.
  • Expect compatibility issues when the model uses unsupported or custom operations.
  • Consider ONNX or other deployment formats if the real requirement is interoperability rather than TensorFlow specifically.

Course illustration
Course illustration

All Rights Reserved.