Speed Up Predictions with ONNX Runtime

You've spent hours training a great model, only to watch it crawl through predictions at inference time. In production, every millisecond counts — your boss wants faster predictions, your users want faster responses, and your GPU bill is climbing. This lesson teaches you how to use ONNX Runtime for faster prediction, a battle-tested inference engine that can speed up your models by 2–5x without retraining.

The Problem This Lesson Solves

Your trained PyTorch or TensorFlow model is slow because those frameworks store extra computational graphs and use eager execution, which is great for training but wasteful during inference. The pain: you can't easily deploy a research model to production without significant speedups, and your current serving solution may have high latency and high resource usage. ONNX Runtime solves this by optimizing and executing your model in a lightweight, cross-platform runtime that's designed for pure inference.

Core Concept / Mental Model

Think of your trained model as a complex recipe. A framework like PyTorch is like a professional kitchen with all the tools, ingredients, and backup staff — great for cooking (training) but heavy for just serving one dish repeatedly. ONNX Runtime is like a streamlined takeout counter: it takes your recipe once, preps everything ahead of time (graph optimizations), and executes it with minimal overhead every time you order (predict).

Key terms:

• ONNX (Open Neural Network Exchange): an open format for representing trained models, independent of the original framework.
• ONNX Runtime: a cross-platform inference engine that loads ONNX models and runs them with optimizations like operator fusion, constant folding, and multi-threading.
• Graph optimization: transforming the model's computation graph for faster execution.

How It Works Step by Step

1. Export your model to ONNX format from PyTorch or TensorFlow using built-in export functions.
2. Optional but recommended: simplify and optimize the ONNX graph using onnxsim or onnxoptimizer.
3. Load the ONNX model with onnxruntime.InferenceSession.
4. Run inference with the session, passing input data as numpy or ORT values.
5. Measure and compare latency and throughput against your original framework.

Why does this work so well? ONNX Runtime applies optimizations like operator fusion (combining multiple ops into one), constant folding (precomputing static values), and CPU/GPU execution providers to squeeze out maximum performance. It also uses multiple threads to parallelize computation.

Hands-On Walkthrough

Step 1: Install Dependencies

pip install onnx onnxruntime torch --quiet

Step 2: Export a PyTorch Model to ONNX

Here's a complete example: train a tiny model, export it, and run inference with ONNX Runtime.

import torch
import torch.nn as nn
import numpy as np

# Define a simple model
class TinyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 2)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(2, 1)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = TinyModel()
model.eval()

# Dummy input matching the input shape
x = torch.randn(1, 10)

# Export to ONNX
torch.onnx.export(
    model,
    x,
    "tiny_model.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
)
print("Model exported successfully")

Step 3: Run Inference with ONNX Runtime

import onnxruntime as ort
import numpy as np

# Create ONNX Runtime session (CPU or CUDA)
so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession("tiny_model.onnx", sess_options=so, providers=["CPUExecutionProvider"])

# Prepare input data
input_data = np.random.rand(1, 10).astype(np.float32)

# Run prediction
outputs = session.run(None, {"input": input_data})
print("Prediction:", outputs[0])

Expected output (values will vary):

Model exported successfully
Prediction: [[0.0234567]]

Step 4: Benchmark Latency

import time

# Warm up the session
for _ in range(10):
    session.run(None, {"input": input_data})

# Measure average latency over 100 runs
start = time.perf_counter()
for _ in range(100):
    session.run(None, {"input": input_data})
end = time.perf_counter()
avg_ort = (end - start) / 100

# Compare with PyTorch eager (no gradient)
with torch.no_grad():
    for _ in range(10):
        model(x)
    start = time.perf_counter()
    for _ in range(100):
        model(x)
    end = time.perf_counter()
    avg_torch = (end - start) / 100

print(f"ONNX Runtime avg latency: {avg_ort*1000:.2f} ms")
print(f"PyTorch avg latency: {avg_torch*1000:.2f} ms")
print(f"Speedup: {avg_torch / avg_ort:.2f}x")

Compare Options / When to Choose What

Pro tip: For most production web APIs, ONNX Runtime is the sweet spot — it's optimized, supported by major cloud providers, and can run on CPU efficiently, saving GPU costs.

Variations

• Use onnxruntime-gpu for CUDA acceleration.
• Convert a TensorFlow model using tf2onnx.
• Use onnxsim to remove unnecessary nodes and improve compatibility.

Troubleshooting & Edge Cases

• Export failure due to unsupported ops: Use a recent ONNX opset version (opset_version=17 or higher) and simplify the model with onnxsim. If an op is still unsupported, check ONNX Runtime's documentation or replace that part of the model.
• Dynamic input shapes: ONNX Runtime expects fixed shapes by default. To handle dynamic shapes, set dynamic_axes in torch.onnx.export and use None for that dimension.
• GPU not used: Ensure you've installed onnxruntime-gpu and listed CUDAExecutionProvider in the session's providers list. Verify with ort.get_available_providers().
• Memory leaks: Reuse the InferenceSession instead of creating new sessions per request.
• Wrong dtype: ONNX runtime expects inputs as numpy arrays of correct dtype (e.g., float32). Cast if necessary.

What You Learned & What's Next

You now understand how to use ONNX Runtime for faster prediction: you can export a PyTorch model to ONNX, load it with ONNX Runtime, and achieve significant speedups with minimal code changes. This is a crucial skill for deploying efficient inference services. In the next lesson, you'll learn how to serve these predictions as a REST API with FastAPI, turning your fast model into a production-ready service.

Remember: The key to fast predictions is not just the model architecture — it's also the inference engine you choose. ONNX Runtime gives you production-grade speed with minimal effort.

Now go convert your next model and measure the speedup yourself!