PyTorch Deployment with ONNX

Real-World Scenario

A fraud detection model needs to run inside a Java-based payment processing service with sub-5ms latency. PyTorch eager mode adds ~15ms Python overhead per call. Exporting to ONNX and running with ONNX Runtime drops inference to 1.2ms — a 12x speedup — and removes the Python dependency entirely.

Exporting a Model to ONNX

import torch
import torch.nn as nn
import torch.onnx
from pathlib import Path

# Define a model
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: torch.Tensor) -> torch.Tensor:
        return self.net(x)

model = FraudDetector()
model.eval()   # CRITICAL: always set eval() before export (affects Dropout/BN)

# Create a dummy input with the correct shape and dtype
dummy_input = torch.randn(1, 20)  # (batch_size, features)

# Export to ONNX
output_path = "fraud_detector.onnx"
torch.onnx.export(
    model,
    dummy_input,
    output_path,
    export_params=True,          # include trained weights
    opset_version=17,            # ONNX opset — use latest stable
    input_names=["features"],    # name the inputs for clarity
    output_names=["fraud_prob"],
    dynamic_axes={
        "features":   {0: "batch_size"},  # variable batch dimension
        "fraud_prob": {0: "batch_size"},
    },
)
print(f"Model exported to {output_path}")

# Verify the exported model
import onnx
onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)
print("ONNX model is valid")
print(f"Graph inputs:  {[i.name for i in onnx_model.graph.input]}")
print(f"Graph outputs: {[o.name for o in onnx_model.graph.output]}")

Running Inference with ONNX Runtime

import onnxruntime as ort
import numpy as np
import time

# Load the ONNX model
session = ort.InferenceSession(
    "fraud_detector.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],  # GPU first, fallback CPU
)

# Print runtime info
print(f"Providers available: {ort.get_available_providers()}")
print(f"Providers in use:    {session.get_providers()}")
print(f"Input names:         {[i.name for i in session.get_inputs()]}")
print(f"Input shapes:        {[i.shape for i in session.get_inputs()]}")

# Single prediction
features = np.random.randn(1, 20).astype(np.float32)  # ONNX Runtime needs float32
result = session.run(
    output_names=["fraud_prob"],
    input_feed={"features": features},
)
fraud_prob = result[0][0, 0]
print(f"Fraud probability: {fraud_prob:.4f}")

# Batch prediction
batch = np.random.randn(1000, 20).astype(np.float32)
t0 = time.perf_counter()
result = session.run(None, {"features": batch})
latency = (time.perf_counter() - t0) * 1000
print(f"Batch of 1000: {latency:.1f}ms ({latency/1000:.3f}ms per sample)")

# Compare with PyTorch eager mode
import torch
torch_model = FraudDetector()   # same model, not loaded from ONNX
torch_model.eval()
torch_batch = torch.tensor(batch)
t0 = time.perf_counter()
with torch.no_grad():
    torch_result = torch_model(torch_batch)
torch_latency = (time.perf_counter() - t0) * 1000
print(f"PyTorch eager:   {torch_latency:.1f}ms")
print(f"ONNX Runtime:    {latency:.1f}ms")
print(f"Speedup:         {torch_latency/latency:.1f}x")

TorchScript Export

import torch
import torch.nn as nn

class SentimentClassifier(nn.Module):
    def __init__(self, vocab_size: int, embed_dim: int, hidden: int, n_classes: int):
        super().__init__()
        self.embedding  = nn.Embedding(vocab_size, embed_dim)
        self.lstm       = nn.LSTM(embed_dim, hidden, batch_first=True)
        self.classifier = nn.Linear(hidden, n_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        embedded = self.embedding(x)
        _, (h_n, _) = self.lstm(embedded)
        return self.classifier(h_n[-1])

model = SentimentClassifier(5000, 64, 128, 2)
model.eval()

# Method 1: torch.jit.trace — fastest, records the execution path
# Use when model has no data-dependent control flow
dummy = torch.zeros(1, 20, dtype=torch.long)
traced_model = torch.jit.trace(model, dummy)
traced_model.save("sentiment_traced.pt")

# Method 2: torch.jit.script — compiles the model with type inference
# Use when model has conditional branches on input values
@torch.jit.script
def classify_text(model: torch.jit.ScriptModule, token_ids: torch.Tensor) -> torch.Tensor:
    logits = model(token_ids)
    return torch.softmax(logits, dim=-1)

# Load and run without needing the original class definition
loaded = torch.jit.load("sentiment_traced.pt")
loaded.eval()

test_input = torch.randint(0, 5000, (1, 20))
probs = torch.softmax(loaded(test_input), dim=-1)
print(f"Sentiment probs: positive={probs[0,1]:.3f}, negative={probs[0,0]:.3f}")

Production Inference Class

import onnxruntime as ort
import numpy as np
from pathlib import Path
import threading

class ONNXModelServer:
    """Thread-safe ONNX model server with warmup and metrics."""

    def __init__(self, model_path: str, n_threads: int = 4):
        opts = ort.SessionOptions()
        opts.intra_op_num_threads  = n_threads
        opts.inter_op_num_threads  = n_threads
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

        self.session    = ort.InferenceSession(model_path, sess_options=opts)
        self.input_name = self.session.get_inputs()[0].name
        self._lock      = threading.Lock()
        self._call_count = 0
        self._total_ms   = 0.0

        # Warmup — first call is slow due to JIT compilation
        warmup_input = np.zeros((1, 20), dtype=np.float32)
        for _ in range(3):
            self.session.run(None, {self.input_name: warmup_input})
        print(f"Model warmed up: {model_path}")

    def predict(self, features: np.ndarray) -> np.ndarray:
        """Run inference. Thread-safe."""
        if features.dtype != np.float32:
            features = features.astype(np.float32)
        if features.ndim == 1:
            features = features.reshape(1, -1)

        import time
        t0 = time.perf_counter()

        with self._lock:
            result = self.session.run(None, {self.input_name: features})[0]
            elapsed = (time.perf_counter() - t0) * 1000
            self._call_count += 1
            self._total_ms   += elapsed

        return result

    def stats(self) -> dict:
        return {
            "calls":      self._call_count,
            "avg_latency_ms": round(self._total_ms / max(self._call_count, 1), 2),
        }


# Usage
server = ONNXModelServer("fraud_detector.onnx")

# Single prediction
features = np.random.randn(20).astype(np.float32)
prob = server.predict(features)[0, 0]
print(f"Fraud probability: {prob:.4f}")

# Simulate concurrent requests
import concurrent.futures
requests = [np.random.randn(20).astype(np.float32) for _ in range(100)]

with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
    futures = [executor.submit(server.predict, r) for r in requests]
    results = [f.result()[0, 0] for f in futures]

print(f"Processed 100 requests: {server.stats()}")
print(f"Mean predicted prob: {np.mean(results):.4f}")