Performance and Accuracy Comparison of PyTorch Models Using Torch-TensorRT Acceleration

Photo by GAMERCOMP.RU on Unsplash

Recently, I’ve been exploring ways to accelerate the inference process. While PyTorch and TensorFlow already provide performance optimizations, there’s still room for improvement.

Inference optimization typically comes into play after a model has been fully developed and is ready for large-scale production. If the inference workload isn’t heavy and latency isn’t a bottleneck, you can often skip this step and simply use the built-in inference methods from PyTorch or TensorFlow.

However, when inference speed becomes critical, converting the model to TensorRT can boost performance by 3x to 5x, an improvement that can be crucial in certain business scenarios.

This post will focus on how to use the Torch-TensorRT library to convert a PyTorch model into a TensorRT-compatible model.

TensorRT SDK vs. torch-tensorrt Library

If you’re confused by these two terms, you’re not alone. I was too when I first came across them. Simply put, TensorRT is part of NVIDIA’s solution ecosystem, while Torch-TensorRT belongs to the PyTorch ecosystem.

To understand TensorRT, we need to mention CUDA programming. CUDA refers to writing GPU-accelerated code directly using NVIDIA’s CUDA toolkit. TensorRT is built on top of CUDA and cuDNN, serving as a specialized inference optimizer and runtime for deep learning models. Instead of writing GPU kernels yourself, you can convert a model into a TensorRT-compatible format and let TensorRT automatically select the best GPU kernels and optimize execution for you.

As the name suggests, Torch-TensorRT converts PyTorch models into TensorRT-compatible models — an optimized version of the original PyTorch model that can deliver significantly faster inference compared to standard PyTorch. Keep in mind that Torch-TensorRT is designed specifically for NVIDIA GPUs, so if you plan to optimize inference on other GPUs (e.g., AMD or Intel), this library won’t apply.

Torch-TensorRT applies a range of optimizations, including layer fusion, kernel auto-tuning, and reduced precision (FP16/INT8) via quantization. The good news is that with this library, you can enable these optimizations by simply passing parameters to its functions, without needing to dive deep into CUDA code yourself.

Code Explanation

The entire script can be run in Google Colab. You can find it in this **GitHub repo**. The Colab environment is T4 GPU, PyTorch 2.8.0+cu126.

The following example looks like a standard model inference script. Since we’re using a transformer model (AutoModel), we also need a corresponding tokenizer to embed the text input. Here, we use AutoTokenizer. Specifically, the code calls AutoTokenizer.from_pretrained to load the pre-trained tokenizer model roberta-base, a widely used option. Then, it calls AutoModel.from_pretrained with the same model name (roberta-base) to load the matching transformer model.

The next step is to prepare a block of random text. To test batch inference, we’ll take the first 16 sentences from this text. In real-world inference tasks, you’d typically use a DataLoader to split long text into manageable chunks. But since our goal here is just to experiment with Torch-TensorRT, we’ll keep it simple and stick with the first 16 sentences.

This batch of 16 sentences is then passed into the tokenizer to generate dense vector representations that capture the semantic meaning of each sentence. In this example, the tokenizer uses the parameter max_length=32 to define the dimension of each vector. In real scenarios, max_length is often set to much higher values, 256, 512 or more, depending on the model and use case. Again, since this is just an experiment, we’ll keep things small and simple.

!pip install -U "nvidia_modelopt[hf]"
!pip install torch-tensorrt

import time
import torch
from transformers import AutoTokenizer, AutoModel
import torch_tensorrt

# 1. Load model & tokenizer
model_name = "roberta-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name).eval().cuda()

# 2. Example batch of sentences
text = """ A good story encourages us to turn the next page and read more. We want to find out what happens next and what the main characters do and what they say to each other. 
We may feel excited, sad, afraid, angry or really happy. This is because the experience of reading or listening to a story is much more likely to make us 'feel' that we are part 
of the story, too. Just like in our 'real' lives, we might love or hate different characters in the story. Perhaps we recognise ourselves or others in some of them. Perhaps we 
have similar problems. Because of this natural empathy with the characters, our brains process the reading of stories differently from the way we read factual information. 
Our brains don't always recognise the difference between an imagined situation and a real one so the characters become 'alive' to us. What they say and do is therefore more meaningful. 
This is why the words and structures that relate a story's events, descriptions and conversations are processed in this deeper way. In fact, cultures all around the world have always 
used storytelling to pass knowledge from one generation to another. Our ancestors understood very well that this was the best way to make sure our histories and information about 
how to relate to others and to our world was not only understood, but remembered too. (Notice that the word ‘history’ contains the word ‘story’ – More accurately, the word ‘story’ 
derives from ‘history’.) Encouraging your child to read or listen to stories should therefore help them to learn a second language in a way that is not only fun, but memorable. 
Let's take a quick look at learning vocabulary within a factual text or within a story. Imagine the readers are eight-year-olds interested in animals. In your opinion, are they more 
likely to remember AND want to continue reading the first or second text? """

texts = [item.strip() for item in text.split(".")][:16]  # adjust batch size here
inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=32).to("cuda")

input_ids = inputs["input_ids"].to(torch.int32)
attention_mask = inputs["attention_mask"].to(torch.int32)

The next step is to perform inference using the transformer. By feeding the model the input_ids and attention_mask prepared earlier, we obtain the model’s output, which we store in the variable baseline_outputs.

We also record the inference time in baseline_time. These will serve as a reference for comparing against the performance and outputs of the TensorRT-optimized models later.

# 3. Baseline PyTorch inference
with torch.no_grad():
    start = time.time()
    baseline_outputs = model(input_ids, attention_mask)
    torch.cuda.synchronize()
    end = time.time()
    baseline_time = end - start
    print(f"PyTorch latency: {baseline_time:.4f} sec")

print("Output shape (PyTorch):", baseline_outputs.last_hidden_state.shape)

The next step is to use the torch_tensorrt module to compile the PyTorch model into TensorRT. This part can be a bit tricky. For the inputs parameter, you need to specify the shape for each input. According to the torch_tensorrt documentation, if your inputs always have a fixed batch size and sequence length, you can provide a single static shape instead of specifying min_shape, opt_shape, and max_shape.

If you want the model to support more flexible input sizes, you can provide all three shapes. Just make sure that for each dimension: min_shape < opt_shape < max_shape under each dimension.

A note of caution: although the torch_tensorrt documentation mentions dynamic shapes, specifying different sequence lengths (the second dimension) will currently cause an error. This indicates that Torch-TensorRT currently only supports dynamic batch sizes (the first dimension of the input shape).

The enabled_precisions argument specifies which numerical precision(s) Torch-TensorRT is allowed to use during computation. You can provide multiple data types, and Torch-TensorRT will automatically choose the appropriate one when compiling the engine.

In general, standard operations such as MatMul and Conv will use FP16 (torch.half) for faster computation, while certain normalization operations or operations involving very small values will use FP32 (torch.float32) to maintain numerical stability.

# Convert model using Torch-TensorRT with enabled_precision of torch.float16
trt_model_float16 = torch_tensorrt.compile(
    model,
    inputs=[
        torch_tensorrt.Input(min_shape=[1, 32], opt_shape=[8, 32], max_shape=[16, 32], dtype=torch.int32),  # input_ids
        torch_tensorrt.Input(min_shape=[1, 32], opt_shape=[8, 32], max_shape=[16, 32], dtype=torch.int32),  # attention_mask
    ],
    enabled_precisions={torch.float16},
)
trt_inputs = {
    "input_ids": input_ids,
    "attention_mask": attention_mask
}

print("Convert to TensorRT float16.")

# Convert model using Torch-TensorRT with enabled_precision of torch.float32
trt_model_float32 = torch_tensorrt.compile(
    model,
    inputs=[
        torch_tensorrt.Input(min_shape=[1, 32], opt_shape=[8, 32], max_shape=[16, 32], dtype=torch.int32),  # input_ids
        torch_tensorrt.Input(min_shape=[1, 32], opt_shape=[8, 32], max_shape=[16, 32], dtype=torch.int32),  # attention_mask
    ],
    enabled_precisions={torch.float32},
)
trt_inputs_float32 = {
    "input_ids": input_ids,
    "attention_mask": attention_mask
}
print("Convert to TensorRT float32.")

When calling trt_model, we need to ensure that the input shapes and data types match those defined in torch_tensorrt.compile. In this example, the trt_inputs variable contains two inputs, each within the allowed shapes and using the torch.int32 data type.

Although TensorRT runtime often prefers int32 inputs for GPU efficiency, many PyTorch or Hugging Face models expect int64 (torch.long) inputs for input_ids and attention_mask. If you choose to keep the original int64inputs, the inputs definition in torch_tensorrt.compile must use the same data type dtype=torch.int64. Be aware that this may slightly slow down inference.

The torch.cuda.synchronize statement is important because it tells the CPU to wait until all previously queued GPU operations have completed. By calling it, we ensure that trt_model(**trt_inputs) has actually finished running on the GPU. Without this, the code would continue executing immediately, and the recorded end time would be much shorter than the actual GPU execution time.

# run model trt_model_float16
start = time.time()
trt_outputs_float16 = trt_model_float16(**trt_inputs)
torch.cuda.synchronize()
end = time.time()
trt_time_float16 = end - start
print(f"Torch-TensorRT latency: {trt_time_float16:.4f} sec")
print("Output shape (TensorRT):", trt_outputs_float16.last_hidden_state.shape)

# run model trt_model_float32
start = time.time()
trt_outputs_float32 = trt_model_float32(**trt_inputs_float32)
torch.cuda.synchronize()
end = time.time()
trt_time_float32 = end - start
print(f"Torch-TensorRT latency: {trt_time_float32:.4f} sec")

print("Output shape (TensorRT):", trt_outputs_float32.last_hidden_state.shape)

Performance Comparison

Since enabled_precisions includes torch.float16 or torch.float32, the inference results may differ numerically from the baseline PyTorch model. The following code demonstrates how to compare the output values between different TensorRT models and the baseline PyTorch model.

Comparing inference execution times is straightforward by looking at baseline_time, trt_time_float16, and trt_time_float32.

import torch
threshold = 0.01

diff_tensor = torch.abs(baseline_outputs.last_hidden_state - trt_outputs_float16.last_hidden_state)
max_diff_float16 = diff_tensor.max().item()
min_diff_float16 = diff_tensor.min().item()
percent_over_threshold_float16 = (diff_tensor > threshold).float().mean().item() * 100
print(f"Max absolute difference: {max_diff_float16}")
print(f"Min absolute difference: {min_diff_float16}")
print(f"Percentage of elements > {threshold}: {percent_over_threshold_float16:.3f}%")

diff_tensor = torch.abs(baseline_outputs.last_hidden_state - trt_outputs_float32.last_hidden_state)
max_diff_float32 = diff_tensor.max().item()
min_diff_float32 = diff_tensor.min().item()
percent_over_threshold_float32 = (diff_tensor > threshold).float().mean().item() * 100
print(f"Max absolute difference: {max_diff_float32}")
print(f"Min absolute difference: {min_diff_float32}")
print(f"Percentage of elements > {threshold}: {percent_over_threshold_float32:.3f}%")

Performance Data Anlalysis

I compared three scenarios:

• Cond-1: Baseline PyTorch inference
• Cond-2: TensorRT inference with enabled_precisions={torch.float16}
• Cond-3: TensorRT inference with enabled_precisions={torch.float32}

The diagram above shows that when compiling a PyTorch transformer model into TensorRT with torch.float16 precision, the inference time for a batch of 16 sentences drops from 0.043 seconds to 0.009 seconds — roughly a 4× speedup. However, if the precision is kept as torch.float32, there is no noticeable improvement in speed or memory usage. In this case, Torch-TensorRT performs static kernel compilation but still uses full FP32.

In terms of inference output accuracy, compared to the baseline PyTorch model, the TensorRT model with torch.float32 shows almost no difference, while the TensorRT model with torch.float16 has about 0.1% of data points with a difference exceeding 0.01. The KDE plot illustrates the distribution of output differences when using torch.float16. As shown, the majority of differences fall between 0 and 0.0025, which is acceptable for most NLP tasks.

That’s all for inference acceleration using Torch-TensorRT. In the next post, I will cover how to convert a PyTorch model into an ONNX model, then build a TensorRT engine manually, and directly manage dynamic shapes, precision, and GPU memory during inference. Happy learning!