Exporting and deploying models

TorchScript and torch.compile, ONNX export, dynamic quantisation, and packaging a model behind an inference API with correct preprocessing.

TorchScript, compile and ONNX

import torch

model.eval()                                  # always export in eval mode

# 1. tracing: run a sample through and record the ops
example = torch.randn(1, 3, 224, 224)
traced = torch.jit.trace(model, example)
traced.save("model_traced.pt")

# tracing ignores Python control flow, so script the model instead when it branches
scripted = torch.jit.script(model)
scripted.save("model_scripted.pt")

# 2. torch.compile: not an export format, a faster runtime for the same Python model
compiled = torch.compile(model, mode="reduce-overhead")
with torch.no_grad():
    out = compiled(example)                   # first call pays the compile cost

# 3. ONNX: the portable interchange format
torch.onnx.export(
    model, example, "model.onnx",
    input_names=["image"], output_names=["logits"],
    dynamic_axes={"image": {0: "batch"}, "logits": {0: "batch"}},
    opset_version=17,
)

import onnxruntime as ort
import numpy as np
session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
node_input = session.get_inputs()[0].name
onnx_out = session.run(None, {node_input: example.numpy()})[0]

with torch.no_grad():
    torch_out = model(example).numpy()
print(np.abs(onnx_out - torch_out).max())     # should be tiny

• Tracing records one execution path. Any if that depends on tensor values is baked in as the branch taken by the example input, which is a silent correctness bug.
• torch.jit.script compiles the Python source and honours control flow, but only supports a restricted subset of Python.
• torch.compile is a runtime optimisation, not a deployment format: your model still needs PyTorch at serving time.
• Export ONNX with dynamic axes for the batch dimension, otherwise the exported graph is fixed to the batch size you traced with.

Quantisation

import torch
import torch.quantization as tq

# dynamic quantisation: weights int8, activations quantised at runtime.
# One of the few things that speeds up CPU inference substantially.
q_model = torch.quantization.quantize_dynamic(
    model.cpu().eval(), {torch.nn.Linear, torch.nn.LSTM}, dtype=torch.qint8)

def size_mb(m):
    buf = torch.save(m.state_dict(), "/tmp/m.pt")
    import os
    return os.path.getsize("/tmp/m.pt") / 1e6

print(size_mb(model), "->", size_mb(q_model))

# static quantisation: calibrate with real data, then convert
model.qconfig = tq.get_default_qconfig("fbgemm")
tq.prepare(model, inplace=True)
with torch.no_grad():
    for i, (x, _) in enumerate(calib_loader):
        if i >= 20:
            break
        model(x)                              # observers collect activation ranges
tq.convert(model, inplace=True)
print(model)

Packaging behind an inference API

import io
import torch
from fastapi import FastAPI, File, UploadFile
from PIL import Image
from torchvision import transforms

app = FastAPI()
model = torch.jit.load("model_traced.pt").eval()

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
CLASSES = ["cat", "dog", "bird", "fish"]

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
    batch = preprocess(image).unsqueeze(0)

    with torch.inference_mode():               # cheaper than no_grad for serving
        logits = model(batch)
        probs = torch.softmax(logits, dim=1)[0]

    top = int(probs.argmax())
    return {"label": CLASSES[top], "confidence": round(float(probs[top]), 4)}

• Keep preprocessing next to the model, versioned together. A serving image pipeline that differs from the training one is the most common cause of a model that scores well offline and badly in production.
• torch.inference_mode() disables version-counter tracking, which is faster than no_grad and is the right choice for serving.
• Warm up with one dummy request before the server accepts traffic so the first user does not pay the load and lazy-initialisation cost.
• Log the input hash and the model version with each prediction. When a score drifts, you need to know which artefact answered.

FAQ

Related

Mixed precision and GPU performance Experiment tracking and reproducibility

Last refreshed 2026-09-18.