Optimize Inference with ONNX

Your PyTorch or Transformers model nails accuracy in training, but in production it’s a different story — latency spikes, GPU memory pressure, and a stack of dependencies that make your model brittle. Every millisecond counts when you’re serving predictions, and your current pipeline just isn’t cutting it. That’s where optimizing inference with ONNX becomes your escape hatch: a vendor-neutral format that turns your model into a lean, fast runtime artifact, ready to deploy anywhere from a serverless function to a Jetson device. This lesson is your hands-on playbook to cut inference latency, shrink memory, and make your models actually ship-ready.

The problem this lesson solves

In production, the metrics that matter shift from training accuracy to inference efficiency. But most ML engineers hit the same walls:

• Latency is too high — your model takes 200ms per prediction when the SLO demands under 50ms.
• Dependencies are heavy — PyTorch alone can weigh 2GB+; your Docker image is a joke.
• GPU/CPU utilization is poor — you’re seeing 20% GPU use and high cost per request.
• Deployment targets vary — the same model must run on both a beefy cloud GPU and an edge device with 512MB RAM.

The root cause: you’re running a training framework at inference time. Frameworks like PyTorch and TensorFlow are optimized for flexible compute graphs — great for backprop, terrible for fast forward passes. They carry graph evaluation overhead, dynamic shapes, and debugging hooks you’ll never use.

ONNX solves this by letting you export your model to a static, optimized compute graph that dedicated inference runtimes like ONNX Runtime (ORT) can JIT-compile, fuse operators, and quantize with minimal effort.

Core concept / mental model

Think of your trained model as a raw recipe. PyTorch is the full kitchen where you wrote it — you can cook there, but you’re also paying to maintain the fryers, mixers, and the chef’s table. ONNX is the prepped meal: all ingredients chopped, measured, and packaged into a standardized tray that any microwave (runtime) can heat up instantly.

More formally:

• ONNX (Open Neural Network Exchange) is an open graph format that captures the model’s architecture, weights, and ops in a neutral schema.
• ONNX Runtime is a cross-platform inference engine that loads the .onnx file and executes it with aggressive optimizations: operator fusion, graph simplification, and kernel selection for your hardware.

A simple mental model of the pipeline:

training framework (PyTorch) → export to ONNX (.onnx) → ONNX Runtime inference (C++/CPU/GPU)

The export step bakes the model’s math into a static graph. The runtime then applies hardware-specific optimizations that are impossible in eager mode.

How it works step by step

To get from your PyTorch model to blazing-fast ONNX inference, follow these five steps:

1. Prepare your model for export

Your model must be in eval mode and all parameters frozen (no gradients). Dynamic operations like if branches on tensor values may fail — ONNX wants a static graph, so standardize input shapes if possible, or use dynamic axes as a fallback (more on that later).

2. Export to ONNX with torch.onnx.export

This API traces your model by running a dummy input through it and recording the graph of tensor operations. You specify:

• input_names / output_names — used with dynamic axes.
• opset version — the ONNX operator version; higher means newer ops, but older runtimes may not support them.
• dynamic_axes — to allow variable batch size or sequence length.

3. Validate the exported graph

Use ONNX Checker to catch structural issues and onnx.shape_inference to verify tensor shapes.

4. Run inference with ONNX Runtime

Create an InferenceSession, bind to CPU/GPU, and optionally enable graph optimizations at session creation.

5. Quantize (optional)

Apply dynamic or static quantization to shrink the model and speed up CPU inference with minimal accuracy loss.

Hands-on walkthrough

We’ll optimize inference for a small neural network — a two-layer Linear + ReLU model — and measure the speedup. This mirrors the workflow you’d use for real models like BERT or ResNet.

Step 1: Train and export a PyTorch model

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, in_features=768, hidden=256, out_features=10):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden, out_features)

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

model = SimpleNet()
model.eval()

# Dummy input matching the expected batch size (1 for production)
dummy_input = torch.randn(1, 768)

# Export to ONNX
torch.onnx.export(
    model,
    dummy_input,
    "simple_model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
    opset_version=17
)
print("Model exported to simple_model.onnx")

Step 2: Validate the ONNX graph

import onnx

onnx_model = onnx.load("simple_model.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX model is valid:")
print(onnx.helper.printable_graph(onnx_model.graph)[:200])

Expected output: a message confirming validity and the start of the graph description.

Step 3: Run inference with ONNX Runtime and measure latency

import onnxruntime as ort
import numpy as np
import time

# Choose provider: CPU or CUDA
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if ort.get_device() == 'GPU' else ['CPUExecutionProvider']
session = ort.InferenceSession("simple_model.onnx", providers=providers)

# Prepare input (same shape as with dynamic batch)
input_data = np.random.randn(1, 768).astype(np.float32)

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

# Measure inference
start = time.perf_counter()
for _ in range(100):
    session.run(None, {"input": input_data})
end = time.perf_counter()

latency_ms = (end - start) / 100 * 1000
print(f"Average ONNX inference latency: {latency_ms:.2f} ms per prediction")

Expected output: A latency number — typically 2–10x faster than the original PyTorch model in eager mode.

Step 4: Quantize for CPU speedup

from onnxruntime.quantization import quantize_dynamic, QuantType

quantized_model_path = "simple_model_quantized.onnx"
quantize_dynamic(
    "simple_model.onnx",
    quantized_model_path,
    weight_type=QuantType.QInt8
)

# Infer with quantized model
session_q = ort.InferenceSession(quantized_model_path, providers=['CPUExecutionProvider'])
# Measure as above...
print("Quantized model created and ready")

Expected output: Success message; quantized model often runs 1.5–3x faster on CPU with negligible accuracy drop for this small model.

Compare options / when to choose what

ONNX is not the only optimization path. Compare it with other popular approaches:

When to pick ONNX: - Multi-framework teams (mix PyTorch and TensorFlow) — ONNX unifies packaging. - Multi-platform deployment (cloud GPU + edge ARM) — ONNX Runtime runs on both. - When you want continuous optimization (quantization, graph cutting) without rewriting your model.

When to avoid ONNX: - Only need one hardware target (e.g., NVIDIA data center) — TensorRT may be faster. - Model uses exotic ops not yet in ONNX — export may fail. - Edge case where you need full training-inference parity — ONNX drops some details (like dropout at inference).

Troubleshooting & edge cases

"Export failed: Unsupported operator"

Your model uses custom ops or dynamic control flow. - Fix: Replace dynamic loops with static ones, or use torch.onnx.is_custom_op_supported check. For attention-like ops, export to an ops version that includes them.

"Shape mismatch at runtime"

You didn’t set dynamic axes but pass different batch sizes. - Fix: Re-export with dynamic_axes for the batch size, or clamp input batch to match the graph.

"BatchNorm / Dropout behave oddly at inference"

These layers have training/inference modes. If you forget model.eval() before export, the model may bake training statistics. - Fix: Always call model.eval() and disable gradient computation (with torch.no_grad():).

"Latency still high"

• Check: Are you binding to GPU? providers list matters; if CUDA provider fails, it silently falls back to CPU unless you pass providers argument explicitly.
• Check: Are you doing dynamic shape inference for every call? Set dynamic_axes only where needed; static shapes allow more optimizations.
• Check: Graph optimizations — set optimization_level=ort.GraphOptimizationLevel.ORT_ENABLE_ALL at session creation.

What you learned & what's next

You now have a working toolkit to optimize inference with ONNX. Specifically, you learned:

• How to export a PyTorch model to ONNX — the core of cross-framework portability.
• How to validate and run an ONNX model with ONNX Runtime.
• How to quantize and measure latency to verify improvements.
• How to choose ONNX over alternatives like TensorRT or TorchScript based on your deployment scope.

You’ve also built the mental model that training-time frameworks are not inference tools — ONNX is your production vehicle.

Next step in your Applied AI engineering path: move beyond single-model optimization to orchestrating multiple models in a serving pipeline — think model versioning, batching, and A/B testing. ONNX will be your foundation for that. Keep this .onnx file handy; you’ll feed it into a serving framework soon.

Now it’s your turn: take your own model, export it, and see how much you can shave off latency. Measure twice — optimize forever.