别只装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.1Python 3.8虚拟环境PyTorch 1.10.0和torchvision 0.11.1提示使用conda管理Python环境可以避免依赖冲突conda create -n tensorrt_demo python3.8 conda activate tensorrt_demo安装必要的Python包pip install torch1.10.0 torchvision0.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(pretrainedTrue) 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} } )关键参数说明参数作用推荐值dynamic_axes允许输入输出batch维度动态变化建议至少支持动态batchopset_versionONNX算子集版本11或更高do_constant_folding是否优化常量True常见问题如果遇到Unsupported: ONNX export of operator for training mode错误确保模型处于eval模式model.eval()3. ONNX模型转换与TensorRT优化获得ONNX模型后我们需要通过TensorRT的优化器生成优化后的引擎。这个过程中TensorRT会执行多种优化层融合合并连续的卷积、BN和激活层精度校准FP16或INT8量化内核自动调优选择最适合当前硬件的计算内核转换代码示例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)优化选项对比优化级别速度提升精度损失适用场景FP32基准无最高精度要求FP161.5-3x轻微大多数应用INT83-5x明显对延迟敏感场景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(bindingsbindings) cuda.memcpy_dtoh(output_data, outputs[0])实测性能数据NVIDIA T4 GPU指标PyTorchTensorRT-FP32TensorRT-FP16延迟(ms)15.29.85.3吞吐量(qps)65102189GPU显存(MB)1245893562从数据可以看出即使是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 调试技巧当遇到转换错误时可以尝试使用trtexec命令行工具检查ONNX模型/usr/src/tensorrt/bin/trtexec --onnxresnet18.onnx --verbose简化模型结构逐步排查问题层检查ONNX算子支持情况for i in range(trt.get_plugin_registry().num_plugins): print(trt.get_plugin_registry().get_plugin_creator(i).name)在真实项目中第一次成功转换TensorRT引擎后建议建立自动化测试流程。我通常会创建一个包含以下步骤的CI/CD流水线模型训练完成后自动导出ONNX触发TensorRT转换任务运行精度验证测试性能基准测试生成对比报告这种端到端的自动化流程可以确保每次模型更新都能获得最优的推理性能。