在AI模型部署中，推理延迟和吞吐量是决定用户体验和运营成本的关键因素。对于在NVIDIA GPU上运行的模型，想要获得极致的性能，NVIDIA TensorRT是事实上的标准优化工具。然而，直接使用TensorRT API进行部署往往涉及复杂的序列化和内核定制。本文将介绍一种更高效、更灵活的方法：利用ONNX Runtime (ORT) 的TensorRT Execution Provider (TRT EP)，实现PyTorch模型从训练到高性能部署的平滑过渡，通常可以实现5倍甚至更高的加速。

1. 为什么选择 ONNX Runtime + TensorRT?

• 模型兼容性: ONNX作为中间格式，支持PyTorch、TensorFlow等主流框架，实现模型的可移植性。
• 易用性: ONNX Runtime提供统一的API接口，无论是CPU、CUDA还是TensorRT，只需更改配置即可切换后端。
• 极致性能: TRT EP允许ORT利用TensorRT的图优化、内核融合和量化等能力，将模型编译成高度优化的执行引擎。

2. 环境准备

要使用TRT EP，您需要安装ONNX Runtime的GPU版本，并且确保您的系统配置了CUDA和TensorRT库（ORT安装包通常会捆绑或依赖特定版本的TensorRT）。

******bash

确保安装的是支持CUDA和TensorRT的ONNX Runtime版本

pip install onnxruntime-gpu
pip install torch torchvision onnx

3. 步骤一：将PyTorch模型导出为ONNX格式

我们以一个标准的预训练ResNet18模型为例。关键在于定义正确的输入形状和动态轴（如果需要）。

******python
import torch
import onnx
import numpy as np
from torchvision import models

1. 载入模型

model = models.resnet18(pretrained=True)
model.eval()

2. 定义输入张量（Batch Size=1, 3 channels, 224×224）

x = torch.randn(1, 3, 224, 224, requires_grad=True)

3. 定义输入/输出名称和动态轴（如果支持不同Batch Size）

input_names = [ “input_0” ]
output_names = [ “output_0” ]

4. 导出ONNX模型

torch.onnx.export(
model,
x,
“resnet18.onnx”,
verbose=False,
opset_version=14,
input_names=input_names,
output_names=output_names,
dynamic_axes={‘input_0’: {0: ‘batch_size’}, ‘output_0’: {0: ‘batch_size’}}
)

print(“Model exported to resnet18.onnx successfully.”)

验证ONNX模型

onnx_model = onnx.load(“resnet18.onnx”)
onnx.checker.check_model(onnx_model)
print(“ONNX model check passed.”)

4. 步骤二：建立基准性能 (CUDA Execution Provider)

首先，我们使用标准的CUDA Execution Provider建立一个GPU推理基准。这已经比CPU快得多，但尚未经过TRT的深度优化。

******python
import onnxruntime as ort
import time

model_path = “resnet18.onnx”
input_data = np.random.rand(1, 3, 224, 224).astype(np.float32)

使用CUDA Execution Provider作为基准

session_cuda = ort.InferenceSession(
model_path,
providers=[‘CUDAExecutionProvider’],
# provider_options=[
# {‘device_id’: 0} # 如果有多张卡，可以指定
# ]
)

input_name = session_cuda.get_inputs()[0].name
output_name = session_cuda.get_outputs()[0].name

预热

for _ in range(5):
_ = session_cuda.run([output_name], {input_name: input_data})

基准测试

times_cuda = []
for _ in range(100):
start_time = time.time()
_ = session_cuda.run([output_name], {input_name: input_data})
times_cuda.append(time.time() – start_time)

mean_cuda_latency = np.mean(times_cuda) * 1000
print(f”CUDA EP 平均推理延迟: {mean_cuda_latency:.2f} ms”)

5. 步骤三：启用TensorRT Execution Provider实现极致加速

现在，我们将Execution Provider切换到TensorrtExecutionProvider。首次运行时，ORT会调用TensorRT引擎对ONNX图进行优化、融合，并生成高效的序列化引擎（Engine）。此过程可能需要几秒到几分钟不等。

注意: 启用TRT EP时，通常需要同时指定CUDAExecutionProvider作为备用，以处理TensorRT不支持的操作。

关键配置

我们通过provider_options来控制TRT的行为，例如设置工作空间内存限制（trt_max_workspace_size）。

******python

启用TensorRT Execution Provider

重要：首次运行会进行模型编译，需要较长时间

session_trt = ort.InferenceSession(
model_path,
providers=[‘TensorrtExecutionProvider’, ‘CUDAExecutionProvider’],
provider_options=[
{
‘trt_max_workspace_size’: 2147483648, # 2GB 工作空间内存
‘trt_fp16_enable’: True # 开启FP16半精度优化，进一步加速
},
{}
]
)

input_name_trt = session_trt.get_inputs()[0].name
output_name_trt = session_trt.get_outputs()[0].name

预热（等待TRT Engine编译完成）

print(“Waiting for TensorRT engine compilation and first few inferences…”)
for _ in range(10):
_ = session_trt.run([output_name_trt], {input_name_trt: input_data})

TensorRT 加速测试

times_trt = []
for _ in range(100):
start_time = time.time()
_ = session_trt.run([output_name_trt], {input_name_trt: input_data})
times_trt.append(time.time() – start_time)

mean_trt_latency = np.mean(times_trt) * 1000

print(f”TensorRT EP (FP16) 平均推理延迟: {mean_trt_latency:.2f} ms”)

性能对比

if mean_cuda_latency > 0:
speedup = mean_cuda_latency / mean_trt_latency
print(f”———————————-“)
print(f”性能提升倍数: {speedup:.2f} X”)
print(f”———————————-“)

6. 预期结果与结论

在一个现代的NVIDIA GPU上（如A100或V100），对于ResNet18这类模型，如果CUDA EP的延迟是5ms，启用TRT EP并结合FP16后，延迟可能降至1ms以下。这轻松实现了5倍以上的加速。

关键加速点总结：
1. 内核融合 (Kernel Fusion): TRT将多个小操作合并为一个GPU内核。
2. 自动调优 (Auto-Tuning): TRT针对特定GPU和张量形状选择最优的计算算法。
3. 半精度计算 (FP16): 显著减少内存带宽需求和计算时间（如果GPU支持）。

通过使用ONNX Runtime的TensorRT Execution Provider，我们可以在不修改模型部署代码主体的情况下，获得TensorRT带来的极致性能提升，是AI基础设施优化的首选方案。