Serialization — From state_dict to .pt2

Three serialization formats, three use cases

state_dict (.pth / .pt) — parameters + buffers as a Python OrderedDict. Most flexible. Requires the model class to load. The default for training and research.
torch.export (.pt2) — exported computational graph + weights. Loadable without the original Python class. The modern path for deployment.
TorchScript (.pt) — older graph + weights format. Still widely supported but legacy for new code.
ONNX (.onnx) — cross-framework standard. The right choice when your serving runtime isn't PyTorch (ONNX Runtime, TensorRT, OpenVINO, browser-side via onnxruntime-web).

What each format gives up

state_dict is portable across PyTorch versions and resilient to model refactors, but you need the Python class to instantiate the model. torch.export and TorchScript are self-contained but tied to the PyTorch runtime. ONNX is portable across runtimes but loses framework-specific optimizations.

The decision tree

Sharing for further training? → state_dict.
Deploying via PyTorch runtime? → torch.export (.pt2).
Deploying via ONNX Runtime / TensorRT / browser? → ONNX (via the dynamo path).
Deploying to mobile? → ExecuTorch (covered in a later lesson).
Deploying to Apple? → CoreML (also a later lesson).

Code

state_dict — flexible but Python-class-dependent·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self): super().__init__(); self.fc = nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP()
torch.save(model.state_dict(), '/tmp/tiny.pth')

# To load — must have the class definition
m2 = TinyMLP()
m2.load_state_dict(torch.load('/tmp/tiny.pth', weights_only=True))
m2.eval()

torch.export (.pt2) — self-contained, modern·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self): super().__init__(); self.fc = nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP().eval()
example = torch.randn(1, 10)

exported = torch.export.export(model, (example,))
torch.export.save(exported, '/tmp/tiny.pt2')

# Load WITHOUT the class definition
loaded = torch.export.load('/tmp/tiny.pt2')
y = loaded.module()(example)        # call .module() to get a callable
print(y.shape)                       # torch.Size([1, 4])

ONNX export — for non-PyTorch runtimes·python

import torch

class TinyMLP(torch.nn.Module):
    def __init__(self): super().__init__(); self.fc = torch.nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP().eval()
example = torch.randn(1, 10)

# In PyTorch 2.x, torch.onnx.export defaults to dynamo=True (modern path)
torch.onnx.export(
    model, (example,), '/tmp/tiny.onnx',
    input_names=['x'], output_names=['y'],
    dynamic_axes={'x': {0: 'batch'}, 'y': {0: 'batch'}},
)

# Then run with ONNX Runtime (separately installed)
# pip install onnxruntime
import onnxruntime as ort
sess = ort.InferenceSession('/tmp/tiny.onnx')
out = sess.run(None, {'x': example.numpy()})
print(out[0].shape)                  # (1, 4)

Serialization — From state_dict to .pt2

Three serialization formats, three use cases

What each format gives up

The decision tree

Code

External links

Exercise

Progress

Comments 0