PyTorch模型部署工具对比：TensorRT vs ONNX Runtime

PyTorch模型部署工具对比：TensorRT vs ONNX Runtime

在实际生产环境中，PyTorch模型的部署优化至关重要。本文通过具体案例对比TensorRT和ONNX Runtime在性能上的差异。

环境准备

我们使用ResNet50模型进行测试，输入尺寸为224x224，批量大小为32。

import torch
import torch.onnx
import numpy as np

class ResNet50Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
    
    def forward(self, x):
        return self.model(x)

model = ResNet50Model().eval()
input_tensor = torch.randn(32, 3, 224, 224)

ONNX导出与转换

# 导出ONNX模型
torch.onnx.export(model, input_tensor, "resnet50.onnx", 
                  export_params=True, opset_version=11)

# 使用ONNX Runtime推理
import onnxruntime as ort
ort_session = ort.InferenceSession("resnet50.onnx")
input_name = ort_session.get_inputs()[0].name

TensorRT转换

# 安装tensorrt
# pip install nvidia-tensorrt --index-url https://pypi.nvidia.com

import tensorrt as trt
import torch.nn.functional as F

def build_engine(onnx_path, engine_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)
    
    with open(engine_path, "wb") as f:
        f.write(engine.serialize())

性能测试

通过1000次推理测试，平均延迟如下：

结论

TensorRT在推理速度上提升明显，但需要额外的构建时间。ONNX Runtime在部署便利性上更优，适合快速原型开发。