Module 2 — Exporting from PyTorch
The ResNet18 we fine-tuned on Fashion-MNIST at the end of the PyTorch course lives inside a Python process. It weighs 45 MB on disk as a state dict, and answers in about 4 ms per image on a laptop GPU. Nothing about that state prevents us from serving it — until the serving team asks for something that runs in a Java microservice, or on an Android phone, or without a PyTorch install. This module turns that same model into a .onnx file that anything can load, and dwells on the two arguments — dynamic axes and input names — that decide whether the resulting artifact is production-ready or subtly broken.
torch.onnx.export, the minimal call
The exporter's entry point takes the model, a sample input, and a destination path. Everything else is optional but often decisive.
import torch
import torchvision
model = torchvision.models.resnet18(weights="DEFAULT")
model.fc = torch.nn.Linear(model.fc.in_features, 10) # 10 Fashion-MNIST classes
model.load_state_dict(torch.load("fashion_resnet18.pt", map_location="cpu"))
model.eval()
example = torch.randn(1, 3, 224, 224) # single image, three channels, 224x224
torch.onnx.export(
model,
example,
"resnet.onnx",
opset_version=17,
input_names=["input"],
output_names=["logits"],
dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
)
The exporter runs the model on example under a tracer that records every tensor operation, then writes the recorded graph out as ONNX. The example is not stored inside the file; it is used to determine shapes, dtypes and the exact code path taken through forward. Every choice made by the tracer — which branch of an if, how many iterations of a for — is baked in. Module 9 explains what to do when that is not enough.
model.eval() is not optional
The .eval() call is easy to forget and expensive to leave out. Dropout in training mode randomly zeros activations; BatchNorm computes batch statistics instead of using its running average. Exporting in training mode captures those behaviours: the resulting .onnx produces different outputs on identical inputs, and every diagnostic points at the runtime while the bug lives in the exporter.
Add the call before the export, and never rely on the caller having set the mode. In production pipelines, the export function should assert the model's mode:
def export_onnx(model, example, path):
assert not model.training, "Call model.eval() before exporting to ONNX"
torch.onnx.export(model, example, path, opset_version=17,
input_names=["input"], output_names=["logits"],
dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}})
Choosing the example input
The example must be representative of what the model will see in production, not "any tensor of the right shape". Three properties matter.
The dtype must match the deployment input. Fashion-MNIST images arrive as float32 after preprocessing; exporting with a float64 example bakes double-precision casts throughout the graph and inflates both file size and inference cost. The shape must contain the batch axis, even if the batch is one: an example of shape (3, 224, 224) produces a model that refuses batches of any size. The values rarely matter — random data is fine — with one exception: if forward inspects tensor values with an if, the branch taken during export depends on those values.
The classical mistake for the ResNet is to hand it a batch of zeros. Zeros propagate through convolutions and BatchNorms without triggering anything unusual, but if the model has a debug path that activates when the input norm exceeds a threshold, that path stays out of the exported graph. When in doubt, use torch.randn with the training normalisation applied.
Dynamic axes: the argument that decides throughput
By default, the exporter freezes every dimension of every tensor to the shape it saw in the example. A model exported with example.shape == (1, 3, 224, 224) accepts inputs of exactly (1, 3, 224, 224). Serving must then call the model once per request, and any GPU parallelism that comes from batching is unreachable.
dynamic_axes marks specific dimensions as variable. The dictionary maps a tensor name to a mapping of axis index to a symbolic name:
dynamic_axes = {
"input": {0: "batch"}, # batch dimension of the input
"logits": {0: "batch"}, # batch dimension of the output
}
For a text encoder, both the batch and the sequence length usually vary at inference time. The second running model of this course — a small transformer that classifies sentences — needs:
dynamic_axes = {
"input_ids": {0: "batch", 1: "seq"},
"attention_mask": {0: "batch", 1: "seq"},
"logits": {0: "batch"},
}
The symbolic names are cosmetic: two dimensions marked "batch" are not linked; the runtime treats each as an independent variable. What matters is that the exporter emits dim_param entries in the ONNX graph rather than fixed sizes, which the Netron inspector confirms visually.
Without dynamic_axes, a GPU-served model runs one request at a time. On modern GPUs, moving from batch 1 to batch 32 typically multiplies throughput by 3 to 10x for the same latency budget. Discovering this after the model has shipped often means re-exporting under production pressure — a scenario worth avoiding by declaring dynamic axes at the first export.
Naming inputs and outputs deliberately
If you omit input_names and output_names, the exporter invents them: input.1, output.1, sometimes worse. Those names then appear in every client that consumes the model:
session.run(None, {"input.1": pixels}) # brittle
session.run(None, {"input": pixels}) # deliberate
Rename them at export. "input" and "logits" for a classifier, "input_ids", "attention_mask" and "logits" for a transformer — names that mean something to the serving team. Once a model has shipped under a name, changing it is a breaking API change for every consumer.
The dynamo exporter, the modern route
torch.onnx.export in PyTorch 2 has two backends. The legacy tracer — described so far — has been stable for years and covers most cases. The dynamo exporter, activated by dynamo=True, uses TorchDynamo to symbolically analyse the model instead of tracing it. It handles more Python control flow, produces cleaner graphs on modern architectures, and is where PyTorch is heading:
torch.onnx.export(
model,
(example,),
"resnet.onnx",
opset_version=18,
input_names=["input"],
output_names=["logits"],
dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
dynamo=True,
)
For a ResNet, the two exporters produce equivalent graphs, and the legacy tracer is still the safer default. For a transformer with a KV cache, the dynamo exporter often succeeds where the tracer fails. When one raises, try the other — the diagnostic sometimes moves from "not supported" to "handled fine".
After the export
Two lines close every export session. First, load the file with onnx.load and run onnx.checker.check_model to catch structural issues before touching the runtime:
import onnx
onnx_model = onnx.load("resnet.onnx")
onnx.checker.check_model(onnx_model)
print(f"exported opset {onnx_model.opset_import[0].version}, {len(onnx_model.graph.node)} nodes")
Then run the numerical parity check of module 4. Never trust a .onnx file until the two are green. An export that passes the checker but disagrees numerically with the source is the specific bug this course exists to prevent.
In summary
torch.onnx.exporttraces the model on an example input and serialises the recorded graph; the example must match production dtype and shape, and the model must be ineval()mode.dynamic_axesdeclares which dimensions can vary at inference time; without it, the model is locked to the example's shape and batched serving becomes impossible.input_namesandoutput_namesshould be chosen deliberately, because they are baked into every consumer of the model and cannot be changed silently.- The dynamo exporter (
dynamo=True) is the modern path and handles more control flow than the legacy tracer; try it when the tracer refuses a model, and pin the opset either way.
Next module: exporting from TensorFlow with tf2onnx, and dealing with the NHWC-versus-NCHW convention flip.