TopolearnProduction runtimes need a Python-free, fully captured graph; torch.export is the supported way to get one.
torch.export.export(model, args, dynamic_shapes=...) traces a model ahead of time into an ExportedProgram: a single graph of ATen operators plus the parameters, with no Python left. Unlike torch.compile, it allows no graph breaks: anything untraceable is an error you must fix, which is what makes the result portable.
Shapes are specialised to the example inputs unless you declare them dynamic. Dim("batch") marks a dimension as symbolic, optionally with min/max. Export records the constraints it relies on and raises if your code silently needs a fixed size. Data-dependent control flow must be rewritten, for example with torch.cond.
What you can do with the result:
torch.export.save(ep, "model.pt2")andtorch.export.load(...);ep.module()returns a callable for testing.- Compile ahead of time with AOTInductor into a shared library callable from C++.
- Convert to ONNX:
torch.onnx.export(model, args, dynamo=True)uses export under the hood. - Lower to ExecuTorch for mobile and edge devices.
torch.export replaces TorchScript (torch.jit.script/trace), which is in maintenance mode. Export models in eval() mode so dropout and batch-norm behave as in inference.
Example
import torch
from torch import nn
from torch.export import export, Dim
class Net(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(8, 2)
def forward(self, x):
return self.fc(x).relu()
batch = Dim("batch")
ep = export(Net().eval(), (torch.randn(4, 8),),
dynamic_shapes={"x": {0: batch}}) # batch size may vary
print(ep.graph) # ATen ops: linear, relu
torch.export.save(ep, "net.pt2")
loaded = torch.export.load("net.pt2")
print(loaded.module()(torch.randn(10, 8)).shape) # torch.Size([10, 2])Remember
- Export = full graph, no Python, no graph breaks allowed
- Declare dynamic dimensions with
Dim - Feeds AOTInductor, ONNX (
dynamo=True) and ExecuTorch; supersedes TorchScript
Watch out for
- Exporting in training mode, baking dropout and batch-statistics behaviour into the graph.