推理引擎选型对比：TensorRT vs ONNX Runtime

在大模型推理场景中，选择合适的推理引擎对性能优化至关重要。本文将从实际测试角度，对比TensorRT与ONNX Runtime在推理性能、部署复杂度和适用场景方面的差异。

环境准备

# 安装依赖
pip install tensorrt onnxruntime onnx

测试模型

使用一个简单的ResNet50模型作为测试案例，导出为ONNX格式：

import torch
import torch.nn as nn
import onnx

class ResNet50(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = torchvision.models.resnet50(pretrained=True)
        
    def forward(self, x):
        return self.backbone(x)

model = ResNet50()
model.eval()

torch.onnx.export(model, torch.randn(1, 3, 224, 224), "resnet50.onnx", 
                  input_names=["input"], output_names=["output"])

TensorRT推理测试

import tensorrt as trt
import pycuda.driver as cuda
import numpy as np

def build_engine(onnx_path):
    builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, trt.Logger(trt.Logger.WARNING))
    with open(onnx_path, "rb") as f:
        parser.parse(f.read())
    
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 30
    engine = builder.build_engine(network, config)
    return engine

ONNX Runtime推理测试

import onnxruntime as ort

session = ort.InferenceSession("resnet50.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 测试推理时间
import time
inputs = np.random.randn(1, 3, 224, 224).astype(np.float32)
start = time.time()
outputs = session.run([output_name], {input_name: inputs})
end = time.time()
print(f"ONNX Runtime 推理时间: {end - start:.4f}s")

性能对比结果

总结

TensorRT在性能上显著优于ONNX Runtime，但部署门槛较高，需要构建引擎；而ONNX Runtime更易用，适合快速验证和部署。

本测试基于社区常见场景，实际应用中需根据具体硬件和性能需求做最终决策。