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: records ops with example input (handles arbitrary python via recorded ops)
traced = torch.jit.trace(model, torch.randn(1, 8), check_trace=True)
traced.save("traced.pt")

# Script: compiles the python (needs torch-differentiable subset)
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.