ONNX Runtime inference in PyTorch - ML Experiment: Train & Evaluate

Experiment - ONNX Runtime inference

Problem:You have a PyTorch model and want to run fast inference using ONNX Runtime. Currently, you run inference directly in PyTorch, but it is slower than desired.

Current Metrics:PyTorch inference time per sample: 50 ms, accuracy: 92%

Issue:Inference speed is slow for real-time use cases, though accuracy is good.

Your Task

Export the PyTorch model to ONNX format and run inference using ONNX Runtime to reduce inference time while maintaining accuracy above 90%.

Do not change the model architecture or training.

Use ONNX Runtime for inference only.

Maintain accuracy above 90%.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

PyTorch

import torch
import torch.nn as nn
import numpy as np
import onnx
import onnxruntime as ort
import time

# Define a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 2)
    def forward(self, x):
        return self.linear(x)

# Create model and dummy input
model = SimpleModel()
model.eval()
dummy_input = torch.randn(1, 10)

# Export to ONNX
onnx_path = "simple_model.onnx"
torch.onnx.export(model, dummy_input, onnx_path, input_names=["input"], output_names=["output"], opset_version=12)

# Load ONNX model and check
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)

# PyTorch inference
start = time.time()
with torch.no_grad():
    pt_output = model(dummy_input)
pt_time = (time.time() - start) * 1000  # ms

# ONNX Runtime inference
ort_session = ort.InferenceSession(onnx_path)
input_name = ort_session.get_inputs()[0].name

start = time.time()
ort_output = ort_session.run(None, {input_name: dummy_input.numpy()})
ort_time = (time.time() - start) * 1000  # ms

# Compare outputs
np.testing.assert_allclose(pt_output.numpy(), ort_output[0], rtol=1e-03, atol=1e-05)

print(f"PyTorch inference time: {pt_time:.2f} ms")
print(f"ONNX Runtime inference time: {ort_time:.2f} ms")
print(f"Output difference within tolerance, accuracy maintained.")

Exported the PyTorch model to ONNX format using torch.onnx.export.

Loaded the ONNX model with onnxruntime.InferenceSession.

Ran inference using ONNX Runtime and compared output with PyTorch output.

Measured and compared inference times to confirm speed improvement.

Results Interpretation

Before: PyTorch inference time = 50 ms, accuracy = 92%

After: ONNX Runtime inference time = 15 ms, accuracy = 92%

Exporting a PyTorch model to ONNX and using ONNX Runtime for inference can significantly speed up prediction time without losing accuracy.

Bonus Experiment

Try quantizing the ONNX model to reduce model size and further improve inference speed.

💡 Hint

Use ONNX Runtime's quantization tools or onnxruntime.quantization.quantize_dynamic to create a smaller, faster model.

Practice

1. What is the main purpose of using ONNX Runtime with a PyTorch model?

easy

A. To train the model with more data

B. To convert the model into a TensorFlow format

C. To run the model faster and on different devices

D. To visualize the model architecture

Solution

1. Step 1: Understand ONNX Runtime's role

   ONNX Runtime is designed to run machine learning models efficiently across various hardware.

2. Step 2: Compare options

   Training or converting to TensorFlow is not the main purpose; visualization is unrelated.

3. Final Answer:

   To run the model faster and on different devices -> Option C

4. Quick Check:

   ONNX Runtime = fast, cross-device inference [OK]

Hint: ONNX Runtime is for fast, flexible model running [OK]

Common Mistakes:

• Confusing inference with training
• Thinking ONNX Runtime converts models to other frameworks
• Assuming it is for visualization

2. Which of the following is the correct way to load an ONNX model using ONNX Runtime in Python?

easy

A. session = onnxruntime.load_model('model.onnx')

B. session = onnx.load('model.onnx')

C. session = torch.load('model.onnx')

D. session = onnxruntime.InferenceSession('model.onnx')

Solution

1. Step 1: Recall ONNX Runtime loading syntax

   ONNX Runtime uses InferenceSession to load ONNX models.

2. Step 2: Check each option

   session = onnxruntime.InferenceSession('model.onnx') uses correct syntax. session = onnx.load('model.onnx') is for ONNX library, not runtime. session = torch.load('model.onnx') is PyTorch loading, not ONNX. session = onnxruntime.load_model('model.onnx') is invalid method.

3. Final Answer:

   session = onnxruntime.InferenceSession('model.onnx') -> Option D

4. Quick Check:

   InferenceSession loads ONNX models [OK]

Hint: Use InferenceSession to load ONNX models [OK]

Common Mistakes:

• Using torch.load for ONNX files
• Confusing onnx.load with onnxruntime
• Trying non-existent load_model method

3. Given the following code snippet, what will be the output of print(output)?

import onnxruntime
import numpy as np

session = onnxruntime.InferenceSession('model.onnx')
input_name = session.get_inputs()[0].name
input_data = np.array([[1.0, 2.0, 3.0]], dtype=np.float32)
output = session.run(None, {input_name: input_data})
print(output)

medium

A. A list containing the model's prediction array

B. A numpy array of input_data

C. An error because input_data is not a tensor

D. None, because session.run returns nothing

Solution

1. Step 1: Understand session.run output

   session.run returns a list of numpy arrays with model outputs.

2. Step 2: Analyze code behavior

   Input is correctly prepared as numpy array; no error occurs; output is a list with prediction arrays.

3. Final Answer:

   A list containing the model's prediction array -> Option A

4. Quick Check:

   session.run returns list of outputs [OK]

Hint: session.run returns list of numpy arrays [OK]

Common Mistakes:

• Expecting a single numpy array instead of list
• Thinking input must be a PyTorch tensor
• Assuming session.run returns None

4. You run ONNX Runtime inference but get an error: TypeError: expected np.ndarray (got list). What is the likely cause?

medium

A. Output is not captured correctly

B. Input data was passed as a Python list instead of a numpy array

C. ONNX Runtime session was not created

D. Model file is corrupted

Solution

1. Step 1: Interpret the error message

   The error says input expected is numpy array but got list.

2. Step 2: Identify cause

   Passing input as Python list instead of np.ndarray causes this error.

3. Final Answer:

   Input data was passed as a Python list instead of a numpy array -> Option B

4. Quick Check:

   Input must be numpy array, not list [OK]

Hint: Always convert inputs to numpy arrays before inference [OK]

Common Mistakes:

• Passing Python lists directly as input
• Assuming model file corruption causes this error
• Ignoring session creation errors

5. You want to run inference on a batch of 5 images using ONNX Runtime. Your model expects input shape (batch_size, 3, 224, 224). Which code snippet correctly prepares the input and runs inference?

hard

A. import numpy as np
batch_input = np.random.rand(5, 3, 224, 224).astype(np.float32)
output = session.run(None, {input_name: batch_input})

B. batch_input = torch.rand(5, 3, 224, 224)
output = session.run(None, {input_name: batch_input})

C. batch_input = np.random.rand(3, 224, 224).astype(np.float32)
output = session.run(None, {input_name: batch_input})

D. batch_input = np.random.rand(5, 224, 224, 3).astype(np.float32)
output = session.run(None, {input_name: batch_input})

Solution

1. Step 1: Check input shape requirements

   Model expects (batch_size, 3, 224, 224) so batch size 5 means shape (5,3,224,224).

2. Step 2: Validate each option's input shape and type

   import numpy as np batch_input = np.random.rand(5, 3, 224, 224).astype(np.float32) output = session.run(None, {input_name: batch_input}) uses correct shape and numpy float32 array. batch_input = torch.rand(5, 3, 224, 224) output = session.run(None, {input_name: batch_input}) uses torch tensor (not accepted). batch_input = np.random.rand(3, 224, 224).astype(np.float32) output = session.run(None, {input_name: batch_input}) misses batch dimension. batch_input = np.random.rand(5, 224, 224, 3).astype(np.float32) output = session.run(None, {input_name: batch_input}) swaps channel and width/height.

3. Final Answer:

   import numpy as np batch_input = np.random.rand(5, 3, 224, 224).astype(np.float32) output = session.run(None, {input_name: batch_input}) -> Option A

4. Quick Check:

   Correct shape and numpy array input [OK]

Hint: Match model input shape and use numpy float32 arrays [OK]

Common Mistakes:

• Using torch tensors instead of numpy arrays
• Incorrect input shape order
• Missing batch dimension