Module 10 — TorchScript, ONNX and serving

The Fashion-MNIST classifier fine-tuned in module 9 lives in a Python process, needs PyTorch installed, and loses its state the moment the process ends. Production is the opposite of all three: predictions run under a service you did not necessarily write, on machines you may not control, and must not depend on the training environment. This module explains the two artifacts that decouple the model from its training code — TorchScript and ONNX — and connects them to a minimal serving endpoint that returns predictions over HTTP.

Course 37 on model registries and course 40 on serving infrastructure go further. Here we cover the export step itself, because that is where the majority of production bugs originate: not in the serving layer, in the artifact that reaches it.

Why "just pickle the model" is not enough

The obvious first attempt is torch.save(model, "model.pt"). It works, and it fails in production for a specific reason: torch.save pickles the Python class, not the computation. Loading requires the exact same class to be importable in the target environment. Rename a file, refactor a submodule, upgrade PyTorch across a major version — the load fails. Deployment engineers, who are not the ones who trained the model, cannot fix this without pulling in the entire training codebase.

state_dict() from module 8 avoids that trap by saving only tensors, but reloading still requires the class definition. Both artifacts are appropriate for research iteration; neither is a production artifact.

TorchScript and ONNX are two different answers to the same requirement: a serialised computation that runs without the Python source code.

TorchScript: trace versus script

TorchScript is a subset of Python that PyTorch can serialise and run in a standalone runtime — including a C++ one, without any Python interpreter. Two ways lead to it, and choosing the right one is the entire subtlety of this section.

torch.jit.trace runs the model on an example input and records the operations that executed. What it captures is a straight-line list of operations, exactly the ones the example triggered.

import torch
from pathlib import Path

model.eval()
example = torch.randn(1, 3, 224, 224)   # one Fashion-MNIST input, resized
traced = torch.jit.trace(model, example)
traced.save("model_traced.pt")

torch.jit.script statically compiles the model's Python source into TorchScript. It handles if, while, list comprehensions and everything that trace erases.

scripted = torch.jit.script(model)
scripted.save("model_scripted.pt")

The decision rule is unambiguous:

Situation	Use
Forward is a straight chain of layers, no control flow	trace
Forward contains if, while depending on tensor values	script
A custom nn.Module with data-dependent branches	script
A pretrained torchvision model like our ResNet18	trace (works, faster to obtain)

trace silently freezes control flow

This is the trap the module exists to warn against. torch.jit.trace records what happened on the example input. An if that took the true branch on that example is frozen as "always true"; the false branch disappears entirely. The model looks correct because it produced the right answer on the sanity-check input, and then behaves incorrectly on any input that would have taken the other branch. Symptom: the served model returns wrong predictions on a fraction of inputs, apparently at random.

The safe habit: check for control flow, then choose. If the forward is a chain of layers — the ResNet case — trace is fine. As soon as forward inspects a tensor value with if, use script.

ONNX: interoperable, one command away

ONNX (Open Neural Network Exchange) is an industry-standard graph format. Exporting a PyTorch model to ONNX makes it consumable by ONNX Runtime, TensorRT, TensorFlow, mobile inference engines, and cloud endpoints from every major provider. It is the artifact to reach for when the serving stack is not PyTorch.

torch.onnx.export(
    model,
    example,
    "model.onnx",
    input_names=["input"],
    output_names=["logits"],
    dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
    opset_version=17,
)

Two arguments matter enormously.

input_names and output_names are the names by which the serving stack refers to the tensors. Choose them deliberately — "input" and "logits", not the defaults — because they are baked into every client that consumes the model.

dynamic_axes marks which dimensions can vary at inference time. Without it, a model exported with a batch of one accepts only batches of one, and serving must run one request at a time. Declaring the batch axis as dynamic lets the runtime batch requests, which is often a 3-10× throughput multiplier on GPUs.

Numerical parity: the mandatory sanity check

The most treacherous export failure is one where the artifact loads, runs, produces reasonable-looking outputs — and disagrees with the original PyTorch model by a percent here, a class there. Verifying numerical parity is the single non-negotiable step of any export pipeline.

import onnxruntime as ort
import numpy as np

torch_out = model(example).detach().numpy()

sess = ort.InferenceSession("model.onnx")
onnx_out = sess.run(None, {"input": example.numpy()})[0]

diff = np.abs(torch_out - onnx_out).max()
print(f"max absolute difference: {diff:.2e}")
assert diff < 1e-4, "ONNX output diverges from PyTorch"

A difference around $10^{-5}$ or $10^{-6}$ is normal — floating-point rounding differs between backends. A difference above $10^{-3}$ signals a real problem: an operator not supported by the target opset, a layer left in training mode, or eval() forgotten before export.

The same check applies to TorchScript: run the traced or scripted module on the example, compare with the original. Never trust a serialised model that has not been compared numerically to its source.

model.eval() before export

Dropout and BatchNorm behave differently in train and eval mode. Exporting a model in train mode bakes the training behaviour — random dropout, batch statistics — into the artifact. The served model produces different outputs on identical inputs, and every diagnostic points at the serving layer while the bug lives in the training script.

The habit is simple and always safe:

model.eval()
example = torch.randn(1, 3, 224, 224)

Every export in this module starts with that line. Skipping it is a specific class of production bug.

A minimal serving endpoint

FastAPI plus ONNX Runtime deliver a working prediction service in twenty lines.

from fastapi import FastAPI
import onnxruntime as ort
import numpy as np

app = FastAPI()
session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
labels = ["T-shirt", "Trouser", "Pullover", "Dress", "Coat",
          "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]

@app.post("/predict")
def predict(pixels: list[list[list[float]]]):
    x = np.asarray(pixels, dtype=np.float32)[None, ...]     # add batch
    logits = session.run(None, {"input": x})[0][0]
    idx = int(logits.argmax())
    return {"label": labels[idx], "confidence": float(np.exp(logits[idx]) / np.exp(logits).sum())}

providers=["CPUExecutionProvider"] is explicit rather than implicit; on a GPU host, CUDAExecutionProvider runs the same graph an order of magnitude faster with no code change beyond that string. This is why ONNX Runtime is worth learning: the artifact stays the same, the execution provider is a configuration.

Every model registered in production carries its verification numbers

Save alongside every deployed artifact the maximum absolute difference against the PyTorch original, the opset version, and the exact preprocessing pipeline. In six months, when a bug surfaces, those three numbers point straight at the cause. Discovering them then, from scratch, is a full day of work.

In summary

• torch.save(model) pickles the class and requires the source at load time; it is unsuitable for production. TorchScript and ONNX serialise the computation and run without the Python source.
• torch.jit.trace records what happened on the example and silently freezes control flow; use script as soon as the forward pass contains if or while on tensor values.
• ONNX dynamic_axes on the batch dimension is what enables batched inference at serving time; without it, throughput collapses.
• Every export needs model.eval() first and a numerical parity check against the original model afterwards; a max absolute difference above $10^{-3}$ signals a real bug.

Next module: the recap and the 40-question exam that covers the ten modules above.