Model Inference and Performance Benchmarking with ONNX Runtime

Hello! Welcome back to our module on model optimization and deployment.

In our last lesson, you successfully exported a PyTorch Whisper model into the ONNX format. This was a crucial step in preparing the model for production, creating a framework-agnostic and portable asset. However, the export process itself doesn't guarantee better performance. The real benefits are realized when this ONNX model is executed by a specialized, high-performance inference engine.

Today, we will complete the optimization loop. Our learning outcome is to perform inference with the exported model using ONNX Runtime and benchmark its performance. You will learn how to load and run your .onnx file in Python, measure its speed against the original PyTorch model using industry-standard metrics, and understand the architectural reasons behind the performance differences.

1. Introducing ONNX Runtime (ORT)

You've learned that ONNX is a model format. ONNX Runtime (ORT) is the high-performance inference engine that brings that format to life. Developed by Microsoft, it's designed from the ground up to accelerate ONNX models across a vast range of hardware and platforms—from massive cloud servers with GPUs to resource-constrained mobile and web clients.

Key features of ONNX Runtime include:

• Cross-Platform: Runs on Linux, Windows, macOS, Android, and iOS.
• Language Bindings: Provides APIs for Python, C++, C#, Java, and JavaScript.
• Execution Providers (EPs): A pluggable architecture that allows ORT to delegate the computation graph to hardware-specific acceleration libraries like NVIDIA's CUDA and TensorRT, Intel's OpenVINO, or ARM's Compute Library.
• Graph Optimizations: Before execution, ORT applies a series of powerful graph transformations, such as operator fusion, constant folding, and dead code elimination, to optimize the model for speed and efficiency.

To get a deeper appreciation for the design and power of ORT, please watch the following video from Microsoft Research.

The diagram below provides a static view of the architecture discussed in the video, showing how an ONNX model is processed, optimized, and executed via different Execution Providers.

2. Running Inference with an ONNX Model

Now, let's move from theory to practice. The core object for running inference in ONNX Runtime's Python API is the InferenceSession. The process is straightforward:

1. Create an InferenceSession, pointing it to your .onnx model file.
2. Prepare your inputs as a dictionary where keys are the model's input names and values are NumPy arrays.
3. Call the run() method of the session with your input dictionary.
4. Process the results, which are returned as a list of NumPy arrays.

For a complex model like Whisper, which involves an encoder, a decoder, and a beam search algorithm, doing this from scratch can be intricate. The Medium article below provides a fantastic, practical example using an "all-in-one" ONNX model for Whisper, similar to what can be produced by advanced tools like Microsoft's Olive. The principles, however, are identical for the encoder-decoder models we exported with optimum.

Let's adapt that example to run the whisper-tiny.en model you exported in the previous lesson.

Step 1: Install onnxruntime
If you haven't already, install the appropriate package.
For CPU: pip install onnxruntime
For GPU: pip install onnxruntime-gpu

Step 2: Inference Script
The following script demonstrates how to load your exported model and transcribe an audio file. We'll use librosa to load the audio and the transformers.WhisperProcessor to prepare the input features, just as we would for the original PyTorch model.

import librosa
import numpy as np
import onnxruntime as ort
from transformers import WhisperProcessor




# --- 1. Load Processor and Audio ---
# Use the processor saved by optimum during export
processor = WhisperProcessor.from_pretrained("./whisper_tiny_onnx") 



# Load a sample audio file (replace with your own)
audio_sample, sampling_rate = librosa.load("your_audio_file.wav", sr=16000)




# --- 2. Prepare Inputs ---
# The processor converts the audio waveform to mel-spectrograms
input_features = processor(audio_sample, sampling_rate=16000, return_tensors="np").input_features




# The decoder needs to know where to start generating.
# We create a single start token ID. The shape is (batch_size, sequence_length).
decoder_input_ids = np.array([[processor.tokenizer.bos_token_id]], dtype=np.int64)




# --- 3. Create ONNX Runtime Session ---
# Use the appropriate provider for your hardware
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']



# We need two sessions: one for the encoder and one for the decoder
encoder_session = ort.InferenceSession("whisper_tiny_onnx/encoder_model.onnx", providers=providers)
decoder_session = ort.InferenceSession("whisper_tiny_onnx/decoder_model.onnx", providers=providers)




# --- 4. Run Encoder ---
# The encoder processes the audio features
encoder_outputs = encoder_session.run(None, {"input_features": input_features})
encoder_hidden_states = encoder_outputs[0]




# --- 5. Generative Decoding Loop (Greedy Search) ---
# This is a simplified generation loop. Production code uses beam search.
for _ in range(processor.tokenizer.max_length_single_sentence):



    # Run the decoder one step at a time
    decoder_outputs = decoder_session.run(
        None,
        {
            "input_ids": decoder_input_ids,
            "encoder_hidden_states": encoder_hidden_states,
        },
    )
    



    # Get the logits for the next token and find the most likely one
    next_token_logits = decoder_outputs[0][0, -1, :]
    next_token_id = np.argmax(next_token_logits)
    



    # Append the new token to our sequence
    decoder_input_ids = np.append(decoder_input_ids, [[next_token_id]], axis=1).astype(np.int64)




    # If the model predicts the end-of-sequence token, we're done
    if next_token_id == processor.tokenizer.eos_token_id:
        break




# --- 6. Decode the Output ---
transcription = processor.batch_decode(decoder_input_ids, skip_special_tokens=True)
print(f"Transcription: {transcription[0]}")

Note: This script implements a simple greedy decoding search. The optimum library provides a more robust ORTModelForSpeechSeq2Seq class that handles the complex generation logic (like beam search) for you, but this manual implementation clearly shows how the encoder and decoder ONNX models are used in a sequence-to-sequence task.

3. Benchmarking: Defining What to Measure

"Faster" is not a precise metric. To properly benchmark performance, especially for generative models, we need to define specific metrics. Since your goal is to become a researcher, understanding these concepts is vital.

From the article, our key metrics will be:

• End-to-End Latency: The total time from when we start processing the audio until the full transcription is generated. This measures the user-perceived delay for a single request.
• Throughput (Tokens per Second): The number of output tokens generated per second. This is a measure of the model's generation speed, independent of the input processing time.

When measuring latency, simply taking one measurement is unreliable. We often look at the distribution of latencies over many runs and report percentiles (P90, P95, P99) to understand worst-case performance.

4. Practical Benchmarking: A Comparative Script

Now, let's write a script to compare our PyTorch and ONNX models. A fair benchmark requires a controlled environment:

• Warm-up: Run the model a few times before starting the timer. This loads all necessary components into memory and warms up the GPU, preventing one-time initialization costs from polluting the results.
• Iteration: Run the inference many times in a loop and calculate the average latency and throughput.

Here is a Python script to perform this benchmark. It measures the average end-to-end latency.

import time
import torch
import librosa
import numpy as np
import onnxruntime as ort
from transformers import WhisperForConditionalGeneration, WhisperProcessor




# --- Configuration ---
MODEL_ID = "openai/whisper-tiny.en"
ONNX_PATH = "./whisper_tiny_onnx"
AUDIO_FILE = "your_audio_file.wav" # Use a consistent audio file
N_WARMUP = 10
N_RUNS = 50




# --- Load Models and Data ---
processor = WhisperProcessor.from_pretrained(MODEL_ID)
audio_sample, sampling_rate = librosa.load(AUDIO_FILE, sr=16000)
input_features = processor(audio_sample, sampling_rate=16000, return_tensors="pt").input_features




# --- PyTorch Model ---
pt_model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID)
device = "cuda" if torch.cuda.is_available() else "cpu"
pt_model.to(device)
pt_input_features = input_features.to(device)




# --- ONNX Runtime Model ---
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if torch.cuda.is_available() else ['CPUExecutionProvider']
ort_session = ort.InferenceSession(f"{ONNX_PATH}/model.onnx", providers=providers)



# Note: For simplicity, we use the merged model.onnx that `optimum` can create.
# For the separate encoder/decoder, the benchmark logic would be inside the generation loop.
onnx_input = {ort_session.get_inputs()[0].name: input_features.cpu().numpy()}





# === PyTorch Benchmark ===
print("--- Benchmarking PyTorch Model ---")



# Warm-up
for _ in range(N_WARMUP):
    _ = pt_model.generate(pt_input_features)



# Timed runs
torch_latencies = []
for _ in range(N_RUNS):
    start_time = time.perf_counter()
    _ = pt_model.generate(pt_input_features, max_new_tokens=100)
    end_time = time.perf_counter()
    torch_latencies.append(end_time - start_time)
avg_torch_latency = sum(torch_latencies) / len(torch_latencies)
print(f"PyTorch Average Latency: {avg_torch_latency * 1000:.2f} ms")





# === ONNX Runtime Benchmark ===
# For a fair comparison, we use optimum's ORTModel, which handles the generation loop internally.
from optimum.onnxruntime import ORTModelForSpeechSeq2Seq

print("\n--- Benchmarking ONNX Runtime Model ---")
ort_model = ORTModelForSpeechSeq2Seq.from_pretrained(ONNX_PATH, provider=providers[0])




# Warm-up
for _ in range(N_WARMUP):
    _ = ort_model.generate(input_features.cpu())



# Timed runs
onnx_latencies = []
for _ in range(N_RUNS):
    start_time = time.perf_counter()
    _ = ort_model.generate(input_features.cpu(), max_new_tokens=100)
    end_time = time.perf_counter()
    onnx_latencies.append(end_time - start_time)
avg_onnx_latency = sum(onnx_latencies) / len(onnx_latencies)
print(f"ONNX Runtime Average Latency: {avg_onnx_latency * 1000:.2f} ms")




# --- Results ---
speedup = avg_torch_latency / avg_onnx_latency
print(f"\nSpeedup: {speedup:.2f}x")

Note: The script above uses optimum's ORTModelForSpeechSeq2Seq for benchmarking, as it provides a direct, high-level generate method that is equivalent to the PyTorch version, ensuring an apples-to-apples comparison. It abstracts away the manual decoding loop we wrote earlier.

Running this script will give you a concrete measure of the performance gain. As shown in the benchmarking chart from the Medium article, the improvements, especially with quantization (which we covered in the previous lesson), can be substantial—often reducing latency by 70% or more.

Conclusion

In this lesson, you have closed the loop on model deployment optimization. You've not only learned how to execute a model outside of its original training framework but also how to rigorously measure and validate the performance benefits of doing so.

Key Takeaways:

• ONNX Runtime (ORT) is the high-performance engine that executes ONNX models, leveraging Execution Providers for hardware acceleration.
• ORT improves performance through extensive graph optimizations like operator fusion and by using hardware-specific compute kernels.
• The onnxruntime.InferenceSession is the primary Python API for loading and running ONNX models.
• Robust benchmarking requires measuring specific metrics like end-to-end latency and throughput (TPS) over multiple, warmed-up iterations.
• The combination of exporting to ONNX and running with ORT almost always yields significant performance improvements over the native PyTorch execution.

Preview of the Next Lesson:

The low latency you've just achieved is not just a number—it's an enabling technology for new applications. In our next lesson, we will explore one such application: streaming ASR systems. We'll discuss the architectural requirements needed to transcribe audio in real-time as it arrives, where the low-latency inference you benchmarked today becomes a critical building block.