PyTorch Deployment with ONNX

Real-World Scenario

A fraud detection model needs to score every payment inside a service that is not written in Python, with a strict per-request latency budget. Exporting the model to ONNX lets the service load it with ONNX Runtime directly — no Python, no PyTorch — and small single-request calls get noticeably cheaper. The steps are the same for any model: export, verify the exported model against the original, and measure.

The Model

Every example below imports this model. The fixed seed stands in for loading real trained weights. Saved as fraud_model.py:

import torch
import torch.nn as nn

class FraudDetector(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(20, 128), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(128, 64), nn.ReLU(),
            nn.Linear(64, 1), nn.Sigmoid(),
        )

    def forward(self, x):
        return self.net(x)

def trained_model():
    torch.manual_seed(0)                 # stands in for loading real trained weights
    return FraudDetector().eval()

Exporting to ONNX

import os
import onnx
import torch
from fraud_model import trained_model

model = trained_model()                  # eval(): Dropout off before export
example = torch.randn(1, 20)

batch = torch.export.Dim("batch")
torch.onnx.export(
    model, (example,), "fraud_detector.onnx",
    input_names=["features"], output_names=["fraud_prob"],
    dynamic_shapes={"x": {0: batch}},    # batch size may vary; 20 features may not
)

onnx_model = onnx.load("fraud_detector.onnx")
onnx.checker.check_model(onnx_model)
print("opset:", onnx_model.opset_import[0].version)
print("ops:  ", [node.op_type for node in onnx_model.graph.node])
for tensor in (*onnx_model.graph.input, *onnx_model.graph.output):
    dims = [d.dim_param or d.dim_value for d in tensor.type.tensor_type.shape.dim]
    print(tensor.name, dims)

for name in sorted(os.listdir(".")):
    if name.startswith("fraud_detector.onnx"):
        print(f"{name:26s} {os.path.getsize(name):>6,} bytes")

$ python export.py
[torch.onnx] Obtain model graph for `FraudDetector([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `FraudDetector([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decompositions...
[torch.onnx] Run decompositions... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
[torch.onnx] Optimize the ONNX graph...
[torch.onnx] Optimize the ONNX graph... ✅
opset: 20
ops:   ['Gemm', 'Relu', 'Gemm', 'Relu', 'Gemm', 'Sigmoid']
features ['batch', 20]
fraud_prob ['batch', 1]
fraud_detector.onnx        10,664 bytes
fraud_detector.onnx.data   43,520 bytes

In PyTorch 2.14, torch.onnx.export captures the model with torch.export and prints each stage as it goes. dynamic_shapes marks dimension 0 of the argument x (named after forward’s parameter) as variable, which is why the graph’s input is ['batch', 20] — any batch size, exactly 20 features. The exported graph contains no Dropout at all, because the model was in eval() mode. The weights went to a separate fraud_detector.onnx.data file; ship both files together.

Running and Verifying with ONNX Runtime

An exported model is only correct if it gives the same answers as the original, so check that before anything else.

import numpy as np
import onnxruntime as ort
import torch
from fraud_model import trained_model

session = ort.InferenceSession("fraud_detector.onnx", providers=["CPUExecutionProvider"])
print("providers:", session.get_providers())
print("input:    ", [(i.name, i.shape, i.type) for i in session.get_inputs()])

batch = np.random.default_rng(0).standard_normal((1000, 20)).astype(np.float32)
(onnx_out,) = session.run(None, {"features": batch})      # None = every output

with torch.no_grad():
    torch_out = trained_model()(torch.from_numpy(batch)).numpy()

print(onnx_out.shape, onnx_out[:3].ravel())
print("max difference vs PyTorch:", np.abs(onnx_out - torch_out).max())

try:
    session.run(None, {"features": batch.astype(np.float64)})
except Exception as e:
    print(type(e).__name__ + ":", str(e).split(" : ")[-1])

$ python ort_infer.py
providers: ['CPUExecutionProvider']
input:     [('features', ['batch', 20], 'tensor(float)')]
(1000, 1) [0.49687028 0.5018607  0.4913918 ]
max difference vs PyTorch: 5.9604645e-08
InvalidArgument: Unexpected input data type. Actual: (tensor(double)) , expected: (tensor(float))

ONNX Runtime takes NumPy arrays keyed by input name and returns a list of outputs. The largest difference from PyTorch across 1,000 inputs is about 6e-8 — float32 rounding noise, not a bug. The final line shows that ONNX Runtime never converts dtypes for you: a float64 array is rejected outright. Data arriving from JSON or pandas is usually float64, so cast it with .astype(np.float32).

Measuring Latency

import time
import numpy as np
import onnxruntime as ort
import torch
from fraud_model import trained_model

session = ort.InferenceSession("fraud_detector.onnx", providers=["CPUExecutionProvider"])
model = trained_model()

def best_ms(fn, repeats=200):
    times = []
    for _ in range(repeats):
        t0 = time.perf_counter()
        fn()
        times.append(time.perf_counter() - t0)
    return min(times) * 1000                      # best case: least noise

for n in (1, 1000):
    x = np.random.default_rng(0).standard_normal((n, 20)).astype(np.float32)
    xt = torch.from_numpy(x)
    with torch.inference_mode():
        eager = best_ms(lambda: model(xt))
    onnx = best_ms(lambda: session.run(None, {"features": x}))
    print(f"batch {n:4d}: PyTorch {eager:.3f} ms   ONNX Runtime {onnx:.3f} ms   "
          f"ratio {eager / onnx:.1f}x")

$ python latency.py
batch    1: PyTorch 0.039 ms   ONNX Runtime 0.008 ms   ratio 4.7x
batch 1000: PyTorch 0.182 ms   ONNX Runtime 0.191 ms   ratio 1.0x

For a single example, ONNX Runtime was several times faster: at this size the work is dominated by per-call overhead, and PyTorch’s Python-level module dispatch costs more than ONNX Runtime’s C++ call. For a batch of 1,000 the two are roughly even, because both spend the time in the same kind of matrix multiplication; run the script a few times and that ratio wobbles on either side of 1. Taking the best of 200 runs filters out scheduling noise. The lesson is not “ONNX is 5x faster” but “measure at the batch size you will actually serve”.

A Thread-Safe Model Server

import threading
from concurrent.futures import ThreadPoolExecutor
import numpy as np
import onnxruntime as ort

class ONNXModelServer:
    def __init__(self, path, n_threads=2):
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = n_threads
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.session = ort.InferenceSession(path, sess_options=opts)
        self.input_name = self.session.get_inputs()[0].name
        self.calls, self._lock = 0, threading.Lock()
        self.predict(np.zeros(20, dtype=np.float32))     # warm-up call

    def predict(self, features):
        x = np.asarray(features, dtype=np.float32).reshape(-1, 20)
        (probs,) = self.session.run(None, {self.input_name: x})   # safe from many threads
        with self._lock:                                  # only our counter needs a lock
            self.calls += 1
        return probs[:, 0]

server = ONNXModelServer("fraud_detector.onnx")
requests = np.random.default_rng(1).standard_normal((100, 20))   # float64, like JSON input
with ThreadPoolExecutor(max_workers=8) as pool:
    results = list(pool.map(server.predict, requests))
print(f"calls: {server.calls}  (100 requests + 1 warm-up)")
print(f"mean probability: {np.mean(results):.4f}")

$ python server.py
calls: 101  (100 requests + 1 warm-up)
mean probability: 0.4947

ONNX Runtime’s session.run is safe to call from many threads at once, so the server only locks around its own counter — wrapping run itself in a lock would serialize every request. predict casts to float32 and reshapes a single row into a batch of one, so callers can pass plain lists or float64 arrays. The warm-up call in __init__ pays one-time initialization costs before real traffic arrives, which is why the counter reads 101.

TorchScript Is Deprecated

Older tutorials export with torch.jit.trace. Besides being deprecated, tracing has a silent failure mode: it records only the path your example input took.

import torch
import torch.nn as nn

class Branchy(nn.Module):
    def forward(self, x):
        if x.sum() > 0:                   # decision depends on the data
            return x * 2
        return -x

import warnings
import torch
from branchy import Branchy

model = Branchy()
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    traced = torch.jit.trace(model, torch.ones(3))     # records the x.sum() > 0 path only
for msg in dict.fromkeys(f"{w.category.__name__}: {str(w.message).split('. ')[0]}." for w in caught):
    print(msg)                                          # each distinct warning once

negative = torch.tensor([-1.0, -2.0, -3.0])
print("eager: ", model(negative))
print("traced:", traced(negative))

$ python trace_trap.py
FutureWarning: `torch.jit.trace` is deprecated.
FutureWarning: `torch.jit.trace_method` is deprecated.
TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect.
eager:  tensor([1., 2., 3.])
traced: tensor([-2., -4., -6.])

The FutureWarnings say TorchScript is on its way out. The TracerWarning is the real problem: the example input ones(3) took the x * 2 branch, so the traced model always doubles, and for a negative input it returns [-2, -4, -6] where the real model returns [1, 2, 3]. No error is raised when the traced model is called.

torch.export: Fail Loudly, Then Fix

torch.export is the replacement. Given the same model, it refuses instead of guessing:

import torch
from branchy import Branchy

try:
    torch.export.export(Branchy(), (torch.ones(3),))
except Exception as e:
    print(type(e).__name__)
    print(str(e).strip().split(".  ")[0])

$ python export_trap.py 2>/dev/null
GuardOnDataDependentSymNode
Could not guard on data-dependent expression Eq(u0, 1) (unhinted: Eq(u0, 1))

2>/dev/null hides the exporter’s diagnostic logging, which is written to stderr; the exception itself is the useful part. x.sum() > 0 can only be known once real data arrives, so there is no single graph to export. Rewriting the branch with torch.cond puts both paths into the graph:

import torch
import torch.nn as nn

class Branchy(nn.Module):
    def forward(self, x):
        return torch.cond(x.sum() > 0, lambda x: x * 2, lambda x: -x, (x,))

program = torch.export.export(Branchy(), (torch.ones(3),))
torch.export.save(program, "branchy.pt2")

loaded = torch.export.load("branchy.pt2").module()   # no Branchy class needed
print(loaded(torch.tensor([1.0, 2.0, 3.0])))
print(loaded(torch.tensor([-1.0, -2.0, -3.0])))

$ python export_cond.py
tensor([2., 4., 6.])
tensor([1., 2., 3.])

torch.cond(predicate, true_fn, false_fn, operands) makes the branch part of the exported program, so both inputs are handled correctly. torch.export.save writes a .pt2 file, and torch.export.load(...).module() gives back a callable module without needing the Branchy class definition — what TorchScript used to provide.

Practice

import numpy as np
import onnxruntime as ort

session = ort.InferenceSession("fraud_detector.onnx")
for shape in [(7, 20), (7, 21)]:
    try:
        (out,) = session.run(None, {"features": np.zeros(shape, dtype=np.float32)})
        print(shape, "->", out.shape)
    except Exception as e:
        print(shape, "->", " ".join(str(e).split(" : ")[-1].split()))

$ python p1_features.py
(7, 20) -> (7, 1)
(7, 21) -> Got invalid dimensions for input: features for the following indices index: 1 Got: 21 Expected: 20 Please fix either the inputs/outputs or the model.

import numpy as np
import onnxruntime as ort
import torch
from fraud_model import trained_model

torch.onnx.export(trained_model(), (torch.randn(1, 20),), "static.onnx",
                  input_names=["features"], verbose=False)     # no dynamic_shapes
session = ort.InferenceSession("static.onnx")
print(session.get_inputs()[0].shape)
try:
    session.run(None, {"features": np.zeros((8, 20), dtype=np.float32)})
except Exception as e:
    print(" ".join(str(e).split(" : ")[-1].split()))

$ python p2_static.py
[1, 20]
Got invalid dimensions for input: features for the following indices index: 0 Got: 8 Expected: 1 Please fix either the inputs/outputs or the model.

import time
import numpy as np
import onnxruntime as ort

session = ort.InferenceSession("fraud_detector.onnx")
requests = np.random.default_rng(0).standard_normal((100, 20)).astype(np.float32)

t0 = time.perf_counter()
one_by_one = [session.run(None, {"features": r[None]})[0] for r in requests]
t1 = time.perf_counter()
(batched,) = session.run(None, {"features": requests})
t2 = time.perf_counter()

print(f"100 calls of 1: {(t1 - t0) * 1000:.2f} ms")
print(f"1 call of 100:  {(t2 - t1) * 1000:.2f} ms")
print("same answers:", np.allclose(np.concatenate(one_by_one), batched, atol=1e-6))

$ python p3_batching.py
100 calls of 1: 1.65 ms
1 call of 100:  0.16 ms
same answers: True