从PyTorch到TensorRT：手把手教你将训练好的模型转成.engine文件（附完整代码）

从PyTorch到TensorRT手把手教你将训练好的模型转成.engine文件附完整代码在深度学习模型部署的实践中TensorRT因其出色的推理性能优化能力而备受青睐。许多开发者在使用PyTorch完成模型训练后都希望将其转换为TensorRT的.engine格式以获得更快的推理速度。本文将详细介绍这一转换过程从ONNX导出到最终生成.engine文件涵盖每一步的操作细节和常见问题解决方案。1. 环境准备与工具安装在开始模型转换前需要确保环境配置正确。TensorRT的安装通常需要与CUDA和cuDNN版本匹配这是保证后续步骤顺利进行的基础。推荐环境配置CUDA 11.xcuDNN 8.xTensorRT 8.xPyTorch 1.10安装TensorRT可以通过NVIDIA官方提供的tar包或deb包进行。以Ubuntu系统为例使用deb包安装的命令如下sudo dpkg -i nv-tensorrt-repo-ubuntu1804-cuda11.4-trt8.2.5.1-ga-20220505_1-1_amd64.deb sudo apt-key add /var/nv-tensorrt-repo-ubuntu1804-cuda11.4-trt8.2.5.1-ga-20220505/7fa2af80.pub sudo apt-get update sudo apt-get install tensorrt验证安装是否成功import tensorrt as trt print(trt.__version__) # 应输出安装的TensorRT版本号注意不同版本的TensorRT对PyTorch和ONNX的支持可能有所不同建议查阅官方文档确认版本兼容性。2. 从PyTorch导出ONNX模型ONNXOpen Neural Network Exchange格式是连接PyTorch和TensorRT的桥梁。将PyTorch模型导出为ONNX格式是转换过程的第一步。导出ONNX的关键参数input_names定义输入张量的名称output_names定义输出张量的名称dynamic_axes指定哪些维度是动态的opset_versionONNX算子集版本一个完整的导出示例import torch import torchvision.models as models # 加载预训练模型 model models.resnet50(pretrainedTrue) model.eval() # 创建虚拟输入 dummy_input torch.randn(1, 3, 224, 224) # 导出ONNX模型 torch.onnx.export( model, dummy_input, resnet50.onnx, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size}, output: {0: batch_size} }, opset_version11 )常见问题及解决方案算子不支持某些PyTorch操作可能没有对应的ONNX实现。解决方法包括使用更高版本的ONNX opset重写模型中使用的不支持算子添加自定义算子实现动态维度处理如果模型需要支持可变输入尺寸务必在dynamic_axes中正确指定。3. ONNX模型验证与优化导出ONNX模型后建议进行验证和优化以确保模型转换的正确性。验证ONNX模型import onnx # 加载ONNX模型 model onnx.load(resnet50.onnx) # 验证模型结构 onnx.checker.check_model(model) # 打印模型信息 print(onnx.helper.printable_graph(model.graph))优化ONNX模型 TensorRT提供了polygraphy工具用于ONNX模型优化polygraphy surgeon sanitize resnet50.onnx -o resnet50_optimized.onnx \ --fold-constants \ --strip优化后的模型通常会移除不必要的节点折叠常量简化计算图减少内存占用4. 使用TensorRT生成.engine文件获得优化后的ONNX模型后就可以使用TensorRT生成最终的.engine文件了。这一过程涉及几个关键决策点精度选择、工作空间大小设置和优化配置。4.1 精度选择TensorRT支持多种精度模式各有优缺点精度模式优点缺点适用场景FP32精度最高速度最慢对精度要求极高的场景FP16速度较快精度损失小需要硬件支持大多数推理场景INT8速度最快需要校准精度损失较大对速度要求极高的场景FP16模式转换示例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(resnet50_optimized.onnx, rb) as f: parser.parse(f.read()) config builder.create_builder_config() config.set_flag(trt.BuilderFlag.FP16) config.max_workspace_size 1 30 # 1GB engine builder.build_engine(network, config) with open(resnet50.engine, wb) as f: f.write(engine.serialize())4.2 INT8量化INT8量化可以显著提升推理速度但需要额外的校准步骤# 创建校准器 class Calibrator(trt.IInt8EntropyCalibrator2): def __init__(self, calibration_data): super().__init__() self.data calibration_data self.current_index 0 def get_batch_size(self): return 1 def get_batch(self, names): if self.current_index len(self.data): batch self.data[self.current_index] self.current_index 1 return [batch.data_ptr()] else: return None def read_calibration_cache(self): return None def write_calibration_cache(self, cache): pass # 使用校准器构建INT8引擎 config.set_flag(trt.BuilderFlag.INT8) config.int8_calibrator Calibrator(calibration_data)4.3 动态形状处理对于需要支持可变输入尺寸的模型需要特别处理动态形状profile builder.create_optimization_profile() profile.set_shape( input, # 输入名称 (1, 3, 224, 224), # 最小形状 (1, 3, 224, 224), # 最优形状 (32, 3, 224, 224) # 最大形状 ) config.add_optimization_profile(profile)5. 性能调优与常见问题解决生成.engine文件后还需要进行性能调优和验证。以下是一些常见问题的解决方案5.1 性能调优技巧工作空间大小适当增加max_workspace_size可以允许TensorRT尝试更多优化策略但会占用更多内存。层融合策略通过config.set_tactic_sources控制TensorRT使用的优化策略。时序缓存保存优化策略缓存可以加速后续引擎构建。config.set_tactic_sources(1 int(trt.TacticSource.CUBLAS) | 1 int(trt.TacticSource.CUDNN)) # 使用时序缓存 cache builder.create_timing_cache(b) config.set_timing_cache(cache, False)5.2 常见错误与解决方案错误类型可能原因解决方案Unsupported ONNX operatorONNX算子不被TensorRT支持使用插件或自定义层实现Invalid dimensions输入形状不匹配检查ONNX模型输入输出形状Out of memory工作空间不足增加max_workspace_size或减小batch sizeFP16 not supported硬件不支持FP16使用FP32或升级硬件5.3 引擎验证生成引擎后建议进行推理测试验证正确性import pycuda.autoinit import pycuda.driver as cuda # 反序列化引擎 with open(resnet50.engine, rb) as f: runtime trt.Runtime(logger) engine runtime.deserialize_cuda_engine(f.read()) # 创建执行上下文 context engine.create_execution_context() # 准备输入输出 input_bind_idx engine.get_binding_index(input) output_bind_idx engine.get_binding_index(output) # 分配GPU内存 input_host cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(input_bind_idx)), dtypenp.float32) output_host cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(output_bind_idx)), dtypenp.float32) input_device cuda.mem_alloc(input_host.nbytes) output_device cuda.mem_alloc(output_host.nbytes) # 执行推理 stream cuda.Stream() cuda.memcpy_htod_async(input_device, input_host, stream) context.execute_async_v2(bindings[int(input_device), int(output_device)], stream_handlestream.handle) cuda.memcpy_dtoh_async(output_host, output_device, stream) stream.synchronize()在实际项目中从PyTorch到TensorRT的转换往往会遇到各种意想不到的问题。例如某些自定义层可能无法直接转换这时就需要编写TensorRT插件来实现相应功能。另外不同版本的TensorRT对ONNX的支持程度也有所不同遇到问题时查阅官方文档和社区讨论往往能快速找到解决方案。