Turbocharge Your AI: Maximizing ONNX Performance on Dedicated Servers

pip uninstall onnxruntime onnxruntime-gpu

pip install onnx onnxruntime

pip install onnx onnxruntime-gpu

import onnxruntime as ort

sess_options = ort.SessionOptions()

# Turn on all optimizations for maximum speed
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

# Tell ORT to save this faster version of the model to a new file
sess_options.optimized_model_filepath = "optimized_model.onnx"

# Run operations sequentially (best for standard deep learning)
sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL

# Set to 0 to let ONNX Runtime automatically detect the best thread count
sess_options.intra_op_num_threads = 0 

# Keep background management threads low to prevent slowdowns
sess_options.inter_op_num_threads = 1

# Tell the software to use the GPU first, with the fastest settings
providers = [
    ('CUDAExecutionProvider', {
        'device_id': 0,                            # Use the main GPU
        'arena_extend_strategy': 'kNextPowerOfTwo',# Manage memory efficiently
        'cudnn_conv_algo_search': 'EXHAUSTIVE'     # Search for the fastest math algorithms
    }),
    'CPUExecutionProvider'
]

# FIRST RUN: Load the original model. 
# The software will optimize it and save it as 'optimized_model.onnx'.
session = ort.InferenceSession("original_model.onnx", sess_options=sess_options, providers=providers)

# FOR FUTURE RUNS: When you restart the server tomorrow, skip the optimization options 
# and load the optimized model directly to save startup time.
# fast_session = ort.InferenceSession("optimized_model.onnx", providers=providers)

import torch
import numpy as np

# 1. Create a tool to bind memory addresses
io_binding = session.io_binding()

# 2. Imagine your data is already sitting on the GPU
input_tensor = torch.randn(1, 3, 224, 224, device='cuda', dtype=torch.float32)

# 3. Link the GPU memory directly to the model input (Zero-Copy)
io_binding.bind_input(
    name='input',
    device_type='cuda',
    device_id=0,
    element_type=np.float32,
    shape=tuple(input_tensor.shape),
    buffer_ptr=input_tensor.data_ptr()
)

# 4. Prepare a space on the GPU for the final answer
output_tensor = torch.empty((1, 1000), device='cuda', dtype=torch.float32)
io_binding.bind_output(
    name='output',
    device_type='cuda',
    device_id=0,
    element_type=np.float32,
    shape=tuple(output_tensor.shape),
    buffer_ptr=output_tensor.data_ptr()
)

# 5. Run the model entirely on the GPU without moving the data
session.run_with_iobinding(io_binding)