别只装TensorRT！用tar包安装后，手把手带你跑通第一个PyTorch模型推理Demo

从PyTorch到TensorRT：手把手实现模型推理加速全流程

刚装好TensorRT却不知道如何验证它的加速效果？本文将带你完整走通PyTorch模型从导出到TensorRT加速的全流程。不同于简单的安装教程，我们聚焦于实际应用场景，通过一个ResNet18分类模型的案例，演示如何利用TensorRT实现3-5倍的推理速度提升。

1. 环境准备与工具链配置

在开始之前，确保你的开发环境满足以下条件：

• Ubuntu 18.04/20.04 LTS（本文以20.04为例）
• NVIDIA驱动≥470版本
• CUDA 11.3和cuDNN 8.2.1
• Python 3.8虚拟环境
• PyTorch 1.10.0和torchvision 0.11.1

提示：使用conda管理Python环境可以避免依赖冲突：

conda create -n tensorrt_demo python=3.8 conda activate tensorrt_demo

安装必要的Python包：

pip install torch==1.10.0 torchvision==0.11.1 tensorrt pycuda onnx onnxruntime

验证TensorRT是否被正确识别：

import tensorrt as trt print(trt.__version__) # 应输出类似8.2.1.8的版本号

2. 构建并导出PyTorch模型

我们从构建一个简单的ResNet18分类模型开始。虽然ResNet18本身不算复杂，但它的结构包含了CNN的典型组件（卷积、BN、ReLU等），非常适合演示优化过程。

模型定义与训练（简化版）：

import torch import torchvision.models as models # 加载预训练模型 model = models.resnet18(pretrained=True) model.eval() # 示例输入 dummy_input = torch.randn(1, 3, 224, 224) # 导出为ONNX格式 torch.onnx.export( model, dummy_input, "resnet18.onnx", input_names=["input"], output_names=["output"], dynamic_axes={ "input": {0: "batch_size"}, "output": {0: "batch_size"} } )

关键参数说明：

常见问题：如果遇到"Unsupported: ONNX export of operator for training mode"错误，确保模型处于eval模式（model.eval()）

3. ONNX模型转换与TensorRT优化

获得ONNX模型后，我们需要通过TensorRT的优化器生成优化后的引擎。这个过程中TensorRT会执行多种优化：

1. 层融合：合并连续的卷积、BN和激活层
2. 精度校准：FP16或INT8量化
3. 内核自动调优：选择最适合当前硬件的计算内核

转换代码示例：

import tensorrt as trt logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, logger) with open("resnet18.onnx", "rb") as f: if not parser.parse(f.read()): for error in range(parser.num_errors): print(parser.get_error(error)) config = builder.create_builder_config() config.set_flag(trt.BuilderFlag.FP16) # 启用FP16加速 config.max_workspace_size = 1 << 30 # 1GB工作空间 serialized_engine = builder.build_serialized_network(network, config) with open("resnet18.engine", "wb") as f: f.write(serialized_engine)

优化选项对比：

4. 性能对比与结果分析

现在我们来对比原始PyTorch模型和TensorRT优化版本的性能差异。我们使用相同的输入数据，分别测试：

• 端到端延迟：包括数据拷贝时间
• 纯计算时间：仅模型推理时间
• 吞吐量：固定时间内能处理的样本数

测试代码关键片段：

# TensorRT推理上下文 runtime = trt.Runtime(logger) with open("resnet18.engine", "rb") as f: engine = runtime.deserialize_cuda_engine(f.read()) context = engine.create_execution_context() # 分配显存 inputs, outputs, bindings = [], [], [] for binding in engine: size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size dtype = trt.nptype(engine.get_binding_dtype(binding)) # 分配输入输出缓冲区 mem = cuda.mem_alloc(size * dtype.itemsize) bindings.append(int(mem)) if engine.binding_is_input(binding): inputs.append(mem) else: outputs.append(mem) # 执行推理 cuda.memcpy_htod(inputs[0], input_data) context.execute_v2(bindings=bindings) cuda.memcpy_dtoh(output_data, outputs[0])

实测性能数据（NVIDIA T4 GPU）：

从数据可以看出，即使是FP32精度，TensorRT也能带来35%的速度提升。而启用FP16后，速度提升达到3倍以上，同时显存占用减少55%。

5. 高级优化技巧与实战建议

5.1 动态形状支持

实际应用中，输入尺寸往往不是固定的。TensorRT支持通过以下方式定义动态形状：

profile = builder.create_optimization_profile() profile.set_shape( "input", # 输入名称 (1, 3, 224, 224), # 最小形状 (8, 3, 224, 224), # 最优形状 (32, 3, 224, 224) # 最大形状 ) config.add_optimization_profile(profile)

5.2 INT8量化

对于极致性能需求，可以考虑INT8量化：

config.set_flag(trt.BuilderFlag.INT8) # 需要提供校准数据 def calibrate(): # 返回校准数据生成器 for _ in range(100): yield [np.random.randn(1, 3, 224, 224).astype(np.float32)] config.int8_calibrator = trt.EntropyCalibrator2(calibrate())

5.3 调试技巧

当遇到转换错误时，可以尝试：

1. 使用trtexec命令行工具检查ONNX模型：

   /usr/src/tensorrt/bin/trtexec --onnx=resnet18.onnx --verbose

2. 简化模型结构，逐步排查问题层

3. 检查ONNX算子支持情况：

   for i in range(trt.get_plugin_registry().num_plugins): print(trt.get_plugin_registry().get_plugin_creator(i).name)

在真实项目中，第一次成功转换TensorRT引擎后，建议建立自动化测试流程。我通常会创建一个包含以下步骤的CI/CD流水线：

1. 模型训练完成后自动导出ONNX
2. 触发TensorRT转换任务
3. 运行精度验证测试
4. 性能基准测试
5. 生成对比报告

这种端到端的自动化流程可以确保每次模型更新都能获得最优的推理性能。