Table of Contents

🚀 Motivation

Having real-time inference is crucial for computer vision applications. In some domains, a 1-second delay in inference could mean life or death.

Imagine sitting in a self-driving car and the car takes one full second to detect an oncoming speeding truck.

Just one second too late, and you could end up in the clouds 👼👼👼

Or if you’re lucky, you get a very up-close view of the pavement.

I hope that shows how crucial real-time inference is.

Thus, having real-time inference capability is paramount and will determine whether a model gets deployed or not. In many cases, you can pick one or the other:

• A fast model with low accuracy
• A slow model with high accuracy

But can we have the best of both worlds? I.e. a fast and accurate model?

That’s what this post is about.

If you’re technical, and this sounds exciting, then let’s dive in! 🏊‍♂️

💻 Installation

Let’s begin with the installation. I will be using a conda environment to install the packages required for this post. Feel free to the environment of your choice.

conda create -n supercharge_timm_tensorrt python=3.11
conda activate supercharge_timm_tensorrt

We’ll be using the timm library to load a pre-trained model and run inference. So let’s install timm.

pip install timm

At the time of writing, there are over 1370 models available in timm. Any of which can be used in this post.

🔧 Load and Infer

Let’s load a top performing model from the timm leaderboard - the eva02_large_patch14_448.mim_m38m_ft_in22k_in1k model.

The plot above shows the accuracy vs inference speed for the EVA02 model.

Look closely, the EVA02 model achieves top ImageNet accuracy (90.05% top-1, 99.06% top-5) but is lags in speed. Check out the model on the timm leaderboard here.

So let’s get the EVA02 model on our local machine

import timm

model_name = 'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k'
model = timm.create_model(model_name, pretrained=True).eval()

Get the data config and transformations for the model

data_config = timm.data.resolve_model_data_config(model)
transforms = timm.data.create_transform(**data_config, is_training=False)

And run an inference to get the top 5 predictions

import torch
from PIL import Image
from urllib.request import urlopen

img = Image.open(urlopen('https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'))

with torch.inference_mode():
    output = model(transforms(img).unsqueeze(0))

top5_probabilities, top5_class_indices = torch.topk(output.softmax(dim=1) * 100, k=5)

Next, decode the predictions into class names as a sanity check

from imagenet_classes import IMAGENET2012_CLASSES

im_classes = list(IMAGENET2012_CLASSES.values())
class_names = [im_classes[i] for i in top5_class_indices[0]]

for name, prob in zip(class_names, top5_probabilities[0]):
    print(f"{name}: {prob:.2f}%")

Top 5 predictions:

Looks like the model is doing it’s job.

Now let’s benchmark the inference latency.

⏱️ Baseline Latency

We will run the inference 10 times and average the inference time.

import time

def run_benchmark(model, device, num_images=10):
    model = model.to(device)
    
    with torch.inference_mode():
        start = time.perf_counter()
        for _ in range(num_images):
            input_tensor = transforms(img).unsqueeze(0).to(device)
            model(input_tensor)
        end = time.perf_counter()
    
    ms_per_image = (end - start) / num_images * 1000
    fps = num_images / (end - start)
    
    print(f"PyTorch model on {device}: {ms_per_image:.3f} ms per image, FPS: {fps:.2f}")

Let’s benchmark on CPU and GPU.

# CPU Benchmark
run_benchmark(model, torch.device("cpu"))

# GPU Benchmark 
if torch.cuda.is_available():
    run_benchmark(model, torch.device("cuda"))

Alright the benchmarks are in

Although the performance on the GPU is not bad, 12+ FPS is still not fast enough for real-time inference. On my reasonably modern CPU, it took over 1.5 seconds to run an inference.

Definitely not self-driving car material 🤷

Now let’s start to improve the inference time.

🔄 Convert to ONNX

ONNX is an open and interoperable format for deep learning models. It lets us deploy models across different frameworks and devices.

The key advantage of using ONNX is that it lets us deploy models across different frameworks and devices, and offers some performance gains.

To convert the model to ONNX format, let’s first install onnx.

pip install onnx

And export the model to ONNX format

import timm
import torch

model = timm.create_model(
    "eva02_large_patch14_448.mim_m38m_ft_in22k_in1k", pretrained=True
).eval()

onnx_filename = "eva02_large_patch14_448.onnx"

torch.onnx.export(
    model,
    torch.randn(1, 3, 448, 448),
    onnx_filename,
    export_params=True,
    opset_version=20,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, 
                  "output": {0: "batch_size"}},
)

If there are no errors, you will end up with a file called eva02_large_patch14_448.onnx in your working directory.

🖥️ ONNX Runtime on CPU

To run the and inference on the ONNX model, we need to install onnxruntime. This is the ’engine’ that will run the ONNX model.

pip install onnxruntime

One (major) benefit of using ONNX Runtime is the ability to run the model without PyTorch as a dependency. This is great for deployment and for running inference in environments where PyTorch is not available.

The ONNX model we exported earlier only includes the model weights and the graph structure. It does not include the pre-processing transforms. To run the inference using onnxruntime, we need to replicate the PyTorch transforms. To find out the transforms that was used, you can print out the transforms.

print(transforms)

Now let’s replicate the transforms using numpy.

def transforms_numpy(image: PIL.Image.Image):
    image = image.convert('RGB')
    image = image.resize((448, 448), Image.BICUBIC)
    img_numpy = np.array(image).astype(np.float32) / 255.0
    img_numpy = img_numpy.transpose(2, 0, 1)
    mean = np.array([0.485, 0.456, 0.406]).reshape(-1, 1, 1)
    std = np.array([0.229, 0.224, 0.225]).reshape(-1, 1, 1)
    img_numpy = (img_numpy - mean) / std
    img_numpy = np.expand_dims(img_numpy, axis=0)
    img_numpy = img_numpy.astype(np.float32)
    return img_numpy

Using the numpy, transforms let’s run inference with ONNX Runtime.

import onnxruntime as ort

# Create ONNX Runtime session with CPU provider
onnx_filename = "eva02_large_patch14_448.onnx"
session = ort.InferenceSession(
    onnx_filename, 
    providers=["CPUExecutionProvider"] # Run on CPU
)

# Get input and output names
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# Run inference
output = session.run([output_name], 
                    {input_name: transforms_numpy(img)})[0]

If we inspect the output shape, we can see that it’s the same as the number of classes in the ImageNet dataset.

Let’s inspect the output.shape:

And printing the top 5 predictions:

We get the same results as the PyTorch model with ONNX Runtime. That’s a good sign!

Now let’s benchmark the inference latency on ONNX Runtime with a CPU provider (backend).

import time

num_images = 10
start = time.perf_counter()
for i in range(num_images):
    output = session.run([output_name], {input_name: transforms_numpy(img)})[0]
end = time.perf_counter()
time_taken = end - start

ms_per_image = time_taken / num_images * 1000
fps = num_images / time_taken

print(f"Onnxruntime CPU: {ms_per_image:.3f} ms per image, FPS: {fps:.2f}")

Ouch! That’s slower than the PyTorch model. What a bummer! It may seem like a step back, but we are only getting started.

Read on.

🖼️ ONNX Runtime on CUDA

Other than the CPU, ONNX Runtime offers other backends for inference. We can easily swap to a different backend by changing the provider. In this case we will use the CUDA backend.

To use the CUDA backend, we need to install the onnxruntime-gpu package.

pip uninstall onnxruntime

pip install onnxruntime-gpu==1.19.2

You can install all the CUDA dependencies using conda with the following command.

conda install -c nvidia cuda=12.2.2 \
                 cuda-tools=12.2.2 \
                 cuda-toolkit=12.2.2 \
                 cuda-version=12.2 \
                 cuda-command-line-tools=12.2.2 \
                 cuda-compiler=12.2.2 \
                 cuda-runtime=12.2.2

Once done, replace the CPU provider with the CUDA provider.

onnx_filename = "eva02_large_patch14_448.onnx"

session = ort.InferenceSession(
    onnx_filename, 
    providers=["CUDAExecutionProvider"] # change the provider 
)

The rest of the code is the same as the CPU inference.

Just with one line of code change, the benchmarks are as follows:

But that’s kinda expected. Running on the GPU, we should expect a speedup.

Failed to load library libonnxruntime_providers_cuda.so 
with error: libcublasLt.so.12: cannot open shared object 
file: No such file or directory

export LD_LIBRARY_PATH="/home/dnth/mambaforge-pypy3/envs/supercharge_timm_tensorrt/lib:$LD_LIBRARY_PATH"

Theres is one more trick we can use to squeeze out more performance - using CuPy for the transforms instead of NumPy.

CuPy is a library that lets us run NumPy code on the GPU. It’s a drop-in replacement for NumPy, so you can just replace numpy with cupy in your code and it will run on the GPU.

Let’s install CuPy compatible with our CUDA version.

pip install cupy-cuda12x

And we can use it to run the transforms.

def transforms_cupy(image: PIL.Image.Image):
    # Convert image to RGB and resize
    image = image.convert("RGB")
    image = image.resize((448, 448), Image.BICUBIC)

    # Convert to CuPy array and normalize
    img_cupy = cp.array(image, dtype=cp.float32) / 255.0
    img_cupy = img_cupy.transpose(2, 0, 1)

    # Apply mean and std normalization
    mean = cp.array([0.485, 0.456, 0.406], dtype=cp.float32).reshape(-1, 1, 1)
    std = cp.array([0.229, 0.224, 0.225], dtype=cp.float32).reshape(-1, 1, 1)
    img_cupy = (img_cupy - mean) / std

    # Add batch dimension
    img_cupy = cp.expand_dims(img_cupy, axis=0)

    return img_cupy

With CuPy, we got a tiny bit of performance improvement:

Using ONNX Runtime with CUDA is a little better than the PyTorch model on the GPU, but still not fast enough for real-time inference.

We have one more trick up our sleeve.

📊 ONNX Runtime on TensorRT

Similar to the CUDA provider, we have the TensorRT provider on ONNX Runtime. This lets us run the model using the TensorRT high performance inference engine by NVIDIA.

From the compatibility matrix, we can see that onnxruntime-gpu==1.19.2 is compatible with TensorRT 10.1.0.

To use the TensorRT provider, you need to have TensorRT installed on your system.

pip install tensorrt==10.1.0 \
            tensorrt-cu12==10.1.0 \
            tensorrt-cu12-bindings==10.1.0 \
            tensorrt-cu12-libs==10.1.0

Next you need to export library path to include the TensorRT library.

export LD_LIBRARY_PATH="/home/dnth/mambaforge-pypy3/envs/supercharge_timm_tensorrt/python3.11/site-packages/tensorrt_libs:$LD_LIBRARY_PATH"

Replace the /home/dnth/mambaforge-pypy3/envs/supercharge_timm_tensorrt/python3.11/site-packages/tensorrt_libs with the path to your TensorRT library.

Otherwise you’ll encounter the following error:

Failed to load library libonnxruntime_providers_tensorrt.so 
with error: libnvinfer.so.10: cannot open shared object file: 
No such file or directory

Next we need so set the TensorRT provider options in ONNX Runtime inference code.

providers = [
    (
        "TensorrtExecutionProvider",
        {
            "device_id": 0,
            "trt_max_workspace_size": 8589934592,
            "trt_fp16_enable": True,
            "trt_engine_cache_enable": True,
            "trt_engine_cache_path": "./trt_cache",
            "trt_force_sequential_engine_build": False,
            "trt_max_partition_iterations": 10000,
            "trt_min_subgraph_size": 1,
            "trt_builder_optimization_level": 5,
            "trt_timing_cache_enable": True,
        },
    ),
]

onnx_filename = "eva02_large_patch14_448.onnx"
session = ort.InferenceSession(onnx_filename, providers=providers)

The rest of the code is the same as the CUDA inference.

Refer to the TensorRT ExecutionProvider documentation for more details on the parameters.

And now let’s run the benchmark:

Running with TensorRT and cupy give us a 4.5x speedup over the PyTorch model on the GPU and 93x speedup over the PyTorch model on the CPU!

Thank you for reading this far. That’s the end of this post.

Or is it?

You could stop here and be happy with the results. After all we already got a 93x speedup over the PyTorch model.

But.. if you’re like me and you wonder how much more performance we can squeeze out of the model, there’s one final trick up our sleeve.

🎂 Bake pre-processing into ONNX

If you recall, we did our pre-processing transforms outside of the ONNX model in CuPy or NumPy.

This incurs some data transfer overhead. We can avoid this overhead by baking the transforms operations into the ONNX model.

Okay so how do we do this?

First, we need to write the preprocessing code as a PyTorch module.

import torch.nn as nn

class Preprocess(nn.Module):
    def __init__(self, input_shape: List[int]):
        super(Preprocess, self).__init__()
        self.input_shape = tuple(input_shape)
        self.mean = torch.tensor([0.4815, 0.4578, 0.4082]).view(1, 3, 1, 1)
        self.std = torch.tensor([0.2686, 0.2613, 0.2758]).view(1, 3, 1, 1)

    def forward(self, x: torch.Tensor):
        x = torch.nn.functional.interpolate(
            input=x,
            size=self.input_shape[2:],
        )
        x = x / 255.0
        x = (x - self.mean) / self.std

        return x

And now export the Preprocess module to ONNX.

input_shape = [1, 3, 448, 448]
output_onnx_file = "preprocessing.onnx"
model = Preprocess(input_shape=input_shape)

torch.onnx.export(
        model,
        torch.randn(input_shape),
        output_onnx_file,
        opset_version=20,
        input_names=["input_rgb"],
        output_names=["output_prep"],
        dynamic_axes={
            "input_rgb": {
                0: "batch_size",
                2: "height",
                3: "width",
            },
        },
    )

Let’s visualize the exported preprocessing.onnx model on Netron.

Next, let’s visualize the original eva02_large_patch14_448 model on Netron.

Note the name of the input node of the eva02_large_patch14_448 model. We will need this for the merge. The name of the input node is input.

Now, we merge the preprocessing.onnx model with the eva02_large_patch14_448 model. To achieve this, we need to merge the output of the preprocessing.onnx model with the input of the eva02_large_patch14_448 model.

To merge the models, we use the compose function from the onnx library.

import onnx

# Load the models
model1 = onnx.load("preprocessing.onnx")
model2 = onnx.load("eva02_large_patch14_448.onnx")

# Merge the models
merged_model = onnx.compose.merge_models(
    model1,
    model2,
    io_map=[("output_preprocessing", "input")],
    prefix1="preprocessing_",
    prefix2="model_",
    doc_string="Merged preprocessing and eva02_large_patch14_448 model",
    producer_name="dickson.neoh@gmail.com using onnx compose",
)

# Save the merged model
onnx.save(merged_model, "merged_model_compose.onnx")

Note the io_map parameter. This lets us map the output of the preprocessing model to the input of the original model. You must ensure that the input and output names of the models are correct.

If there are no errors, you will end up with a file called merged_model_compose.onnx in your working directory. Let’s visualize the merged model on Netron.

Now using this merged model, let’s run the inference benchmark again using the TensorRT provider.

We’ll need to make a small change to how the input tensor is passed to the model.

def read_image(image: Image.Image):
    image = image.convert("RGB")
    img_numpy = np.array(image).astype(np.float32)
    img_numpy = img_numpy.transpose(2, 0, 1)
    img_numpy = np.expand_dims(img_numpy, axis=0)
    return img_numpy

Notice we are no longer doing the resize and normalization inside the function. This is because the merged model already includes these operations.

And the results are in!

That’s a 8x improvement over the original PyTorch model on the GPU and a whopping 123x improvement over the PyTorch model on the CPU! 🚀

Let’s do a final sanity check on the predictions.

Looks like the predictions tally!

🎮 Video Inference

Just for fun, let’s see how fast the merged model runs on a video.

The video inference code is also provided in the repo.

🚧 Conclusion

In this post we have seen how we can supercharge our TIMM models for faster inference using ONNX Runtime and TensorRT.

Or simply check out the Hugging Face Spaces demo below.

Thank you for reading! I hope this has been helpful. If you’d like to find out how to deploy this model on Android check out the following post.

February 7, 2023

PyTorch at the Edge: Deploying Over 964 TIMM Models on Android with TorchScript and Flutter