Skip to content

Export— TorchScript, ONNX, torch.export

Overview

Getting a trained model out of Python has three main doors— each with different compatibility, portability, and tooling support:

  • TorchScript (torch.jit): the classic; script (compiles Python subset) or trace (records ops). Mature but legacy— the compiler is in maintenance mode.
  • ONNX (torch.onnx.export): universal interchange format; goes to TensorRT, ONNX Runtime, CoreML, TVM. Best for interop with non-PyTorch runtimes.
  • torch.export (torch.export.export): the modern path— a strict, serializable graph representation with ExportedProgram, the foundation for AOTInductor and TorchScript replacement.

For new projects: prefer torch.export → AOTInductor (inference) or ONNX (interop). TorchScript remains for legacy serving stacks.

-

torch.export— the Modern Contract

import torch, torch.nn as nn

class Net(nn.Module):
 def __init__(self):
 super().__init__()
 self.fc = nn.Linear(8, 4)
 def forward(self, x):
 return self.fc(x).relu()

model = Net().eval()
ep = torch.export.export(model, (torch.randn(1, 8),)) # freeze + capture
print(type(ep).__name__) # ExportedProgram

# run the exported program (same numerics)
out = ep.module()(torch.randn(1, 8))
print("shape:", out.shape)
torch.export.save(ep, "model.pt") # portable, single file

Why torch.export differs from TorchScript

  • Strict: fails on unsupported constructs instead of silently breaking.
  • Serializable graph: ep.graph_module is a GraphModule you can inspect/transform.
  • AOTInductor: torch._export.aot_compile lowers straight to C++/Triton:
# compile to a single shared object library
lib_path = torch._export.aot_compile(model, (torch.randn(1, 8),))
# ->.so you can dlopen from C++/Python (Ch 07-03)

TorchScript— Script vs Trace

model = Net().eval()

# Trace
traced = torch.jit.trace(model, torch.randn(1, 8), check_trace=True)
traced.save("traced.pt")

# Script
scripted = torch.jit.script(model)
scripted.save("scripted.pt")
Method Pros Cons
trace handles most code bakes input shapes; data-dependent branch issues
script shape-generic, full control flow python subset limits (dicts, dynamic list ops)

TorchScript is legacy; the torch team's focus has moved to torch.export. Use it only for legacy serving.

-

ONNX— Interop Queen

import torch

model = Net().eval()
torch.onnx.export(
 model,
 torch.randn(1, 8),
 "model.onnx",
 input_names=["input"],
 output_names=["output"],
 dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}, # dynamic batch
 opset_version=17,
)

# verify
import onnx
m = onnx.load("model.onnx")
onnx.checker.check_model(m)

ONNX gotchas

  • Dynamic shapes: declare dynamic_axes or shapes are frozen.
  • Ops coverage: some exotic ops export awkwardly— test everything.
  • No autograd: export in eval() mode, torch.no_grad().
  • Quantization: use ONNX Runtime's quantizer or export a PTQ/QAT model.

Which Door for Which Job

Goal Door Why
Serve with Triton/TensorRT ONNX runtime support
Run inside a C++ app torch.export + AOTInductor .so with no runtime dep
Legacy stack (older TorchServe) TorchScript still requested
Highest-perf self-hosted inference torch.export → AOTInductor fused C++/Triton
Prototype/model zoo torch.export strict + inspectable

-

Export Checklist (the discipline)

  1. model.eval() + torch.no_grad().
  2. Pin input dtypes/shapes (declare dynamic axes if needed).
  3. Remove training-only modules (dropout off, BN frozen).
  4. Handle data-dependent control flow: torch.export fails loudly— refactor to masks.
  5. Verify parity: export → load → compare outputs vs eager (within 1e-4).
  6. Keep the original fp32 model around as the reference.
import torch

def check_parity(exported, eager, x, tol=1e-4):
 a = exported(x).detach()
 b = eager(x).detach()
 print("max abs diff:", (a - b).abs().max().item())

-

Key Takeaways

  • Three doors: TorchScript (legacy), ONNX (interop), torch.export (modern, strict).
  • Prefer torch.export for new work; ONNX for ecosystem interop; JIT only for legacy.
  • Export is a contract: strict shapes/dtypes; verify parity every time.
  • AOTInductor turns exported programs into deployable .so files.
  • Keep the fp32 eager model as ground truth after every export.

-