Deployment: torch.export, TorchScript & ONNX

By the end of this lesson you will know the main ways to take a trained PyTorch model out of your notebook and into an app: exporting it with torch.export, converting it to ONNX, and why older code uses TorchScript. You will also shrink a model with quantisation and see what a small prediction service looks like.

Training is only half the job. A model is useful when a website, a phone app or a factory camera can send it an input and get an answer back, quickly and reliably. That step is called deployment. For the wider picture of versioning and monitoring, see MLOps Basics.

What changes when you deploy

During training you care about learning. In deployment you care about four new things:

1. Speed (latency): how long one prediction takes.
2. Size: how much memory and disk the model needs, which matters a lot on phones.
3. Portability: whether it can run without your Python training code, or even without Python.
4. Reliability: the same input must always give the same output, and bad inputs must not crash the service.

Think of a home cook who makes excellent biryani and wants to supply a railway catering service. The recipe is the same, but now it must be packed in standard containers, cooked the same way every time, and delivered on schedule to many trains. Deployment is the packing and delivery.

Before any export: prepare the model

Whatever route you take, start the same way:

1. Load the best weights with load_state_dict (see Saving & Loading Models).
2. Call model.eval(), so dropout is off and batch norm uses its running statistics.
3. Run predictions under torch.inference_mode() (or torch.no_grad()), which skips gradient tracking.
4. Write down the exact preprocessing: image size, normalisation, class names.

Option 1: torch.export

torch.export.export(model, example_inputs) traces your model into an ExportedProgram: a complete, self-contained graph of the operations it performs, without your Python class. It can be saved to a .pt2 file and loaded again elsewhere, and it is the starting point for PyTorch's newer deployment tools, including ExecuTorch for phones and small devices. This is plain code with its real output.

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3)).eval()
x = torch.randn(2, 4)

exported = torch.export.export(model, (x,))       # example inputs as a tuple
print(type(exported).__name__)
torch.export.save(exported, "model.pt2")

loaded = torch.export.load("model.pt2")
print("same output after reload:", torch.allclose(loaded.module()(x), model(x)))

ExportedProgram
same output after reload: True

The program behaves exactly like the original model. By default the export is specialised to the example input's shapes; if your batch size or image size will vary, you declare which dimensions may change with the dynamic_shapes argument.

Option 2: ONNX

ONNX (Open Neural Network Exchange) is an open file format for models, supported by many tools and languages. A model saved as ONNX can run in ONNX Runtime from Python, C++, C#, Java or JavaScript, often faster than plain PyTorch, and without PyTorch installed at all. That makes it popular for serving models inside existing company systems.

PyTorch exports with torch.onnx.export. Since PyTorch 2.9 it uses the torch.export-based exporter by default (dynamo=True). It needs the onnx and onnxscript packages. This example is plain code; it was not run on this machine because those packages were not installed:

import torch

model.eval()
example = torch.randn(1, 3, 224, 224)
onnx_program = torch.onnx.export(model, (example,), dynamo=True)
onnx_program.save("classifier.onnx")

Running it with ONNX Runtime looks like this (also plain code):

import numpy as np
import onnxruntime as ort

session = ort.InferenceSession("classifier.onnx")
input_name = session.get_inputs()[0].name
image = np.random.rand(1, 3, 224, 224).astype(np.float32)   # preprocessed exactly as in training
logits = session.run(None, {input_name: image})[0]
print(logits.argmax(axis=1))

Notice that the second program does not import PyTorch at all. Before trusting an exported model, always compare its outputs with the original PyTorch model on a few real inputs.

TorchScript: the older route

For years, PyTorch's own deployment format was TorchScript, created with torch.jit.trace or torch.jit.script. You will meet it in many existing projects and older tutorials. It is now deprecated: the PyTorch 2.10 release notes say to use torch.export instead, and ExecuTorch instead of TorchScript's mobile runtime. On this machine, with Python 3.14, calling torch.jit.trace printed this warning:

`torch.jit.trace` is not supported in Python 3.14+ and may break. Please switch to `torch.compile` or `torch.export`.

Making models smaller: quantisation

Model weights are normally 32-bit floats: 4 bytes each. Quantisation stores them in fewer bits, most often as 8-bit integers (1 byte each). The model becomes about four times smaller and often faster on CPUs and phones, with a small loss of accuracy.

The idea is simple: find a scale so that the largest weight maps to 127, round every weight to the nearest whole number in that range, and multiply back by the scale when computing. Think of measuring rice with a steel katori instead of a precise scale: you lose a little precision, but it is much quicker and good enough.

Step 1: quantise a layer's weights in NumPy.

import numpy as np

rng = np.random.default_rng(0)
W = rng.normal(0, 0.1, (256, 128)).astype(np.float32)     # one layer's weights
x = rng.normal(size=(8, 256)).astype(np.float32)          # a batch of inputs

scale = np.abs(W).max() / 127                              # largest weight -> 127
W_int8 = np.round(W / scale).astype(np.int8)               # stored as 1-byte integers
W_back = W_int8.astype(np.float32) * scale                 # used when computing

exact = x @ W
approx = x @ W_back
print(f"size float32: {W.nbytes:,} bytes | int8: {W_int8.nbytes:,} bytes")
print(f"largest weight error: {np.abs(W - W_back).max():.5f} (scale {scale:.5f})")
print(f"output error relative to output size: {np.abs(exact - approx).mean() / np.abs(exact).mean():.4f}")

size float32: 131,072 bytes | int8: 32,768 bytes
largest weight error: 0.00177 (scale 0.00354)
output error relative to output size: 0.0103

The weights shrink to a quarter of the size, and the layer's outputs change by only about one percent on average. Real tools refine this with a separate scale per channel and by calibrating on sample data. Always re-measure accuracy on your test set after quantising.

Serving predictions over the web

A common pattern is a small web service: an app sends an image or some numbers, and the service replies with a prediction as JSON. Python frameworks such as FastAPI or Flask make this short. A sketch with FastAPI (plain code):

from fastapi import FastAPI
import torch

app = FastAPI()
model = load_model()                    # build, load_state_dict, eval(): once, at start-up
CLASSES = ["healthy", "leaf blight", "rust", "mosaic"]

@app.post("/predict")
def predict(features: list[float]):
    x = torch.tensor([features], dtype=torch.float32)
    with torch.inference_mode():
        probs = torch.softmax(model(x), dim=1)[0]
    best = int(probs.argmax())
    return {"label": CLASSES[best], "confidence": round(float(probs[best]), 3)}

A reply from such a service looks like this:

{"label": "rust", "confidence": 0.912}

Load the model once when the service starts, not for every request. Validate inputs (length, ranges, file type) before they reach the model. Return the confidence, so the app can ask a human when the model is unsure.

After deployment

Deployment is not the end:

• Monitor inputs and predictions. If the real data starts to look different from the training data (new phone cameras, a new season), accuracy can drop silently. This is called data drift.
• Version every model with its weights, code, preprocessing and test scores, so you can roll back.
• Log errors and slow requests, but protect personal data, as discussed in Responsible Deep Learning.

Common mistakes

• Exporting without model.eval(). Dropout and batch norm behave like training, and predictions are wrong or random.
• Different preprocessing in production. The app resizes or normalises images differently from training, and accuracy collapses without any error.
• Not comparing exported and original outputs. Always check a few inputs give matching results.
• Loading the model on every request. It makes the service slow; load once at start-up.
• Starting new projects on TorchScript. It is deprecated; use torch.export or ONNX.

Key points

• Deployment needs speed, small size, portability and reliability, not only accuracy.
• Always deploy with model.eval(), torch.inference_mode() and the exact training preprocessing.
• torch.export produces a self-contained ExportedProgram (saved as .pt2); TorchScript is deprecated.
• ONNX lets models run in ONNX Runtime without PyTorch, from many languages.
• Quantisation to 8-bit integers makes models about four times smaller with a small accuracy cost; re-test after it.

import numpy as np

rng = np.random.default_rng(0)
W = rng.normal(0, 0.1, (256, 128)).astype(np.float32)
W[:, 0] *= 10                                   # one column with much larger weights
x = rng.normal(size=(8, 256)).astype(np.float32)
exact = x @ W

def rel_error(W_back):
    return np.abs(exact - x @ W_back).mean() / np.abs(exact).mean()

scale = np.abs(W).max() / 127
per_tensor = np.round(W / scale) * scale

col_scale = np.abs(W).max(axis=0, keepdims=True) / 127     # one scale per column
per_channel = np.round(W / col_scale) * col_scale

print(f"one scale for all:   error {rel_error(per_tensor):.4f}")
print(f"one scale per column: error {rel_error(per_channel):.4f}")

one scale for all:   error 0.0751
one scale per column: error 0.0076