从PyTorch到TensorRT：手把手教你将训练好的模型转成.engine文件（附完整代码）
在深度学习模型部署的实践中，TensorRT因其出色的推理性能优化能力而备受青睐。许多开发者在使用PyTorch完成模型训练后，都希望将其转换为TensorRT的.engine格式，以获得更快的推理速度。本文将详细介绍这一转换过程，从ONNX导出到最终生成.engine文件，涵盖每一步的操作细节和常见问题解决方案。
1. 环境准备与工具安装
在开始模型转换前，需要确保环境配置正确。TensorRT的安装通常需要与CUDA和cuDNN版本匹配，这是保证后续步骤顺利进行的基础。
推荐环境配置：
- CUDA 11.x
- cuDNN 8.x
- TensorRT 8.x
- PyTorch 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模型
ONNX（Open Neural Network Exchange）格式是连接PyTorch和TensorRT的桥梁。将PyTorch模型导出为ONNX格式是转换过程的第一步。
导出ONNX的关键参数：
input_names：定义输入张量的名称output_names：定义输出张量的名称dynamic_axes：指定哪些维度是动态的opset_version：ONNX算子集版本
一个完整的导出示例：
import torch import torchvision.models as models # 加载预训练模型 model = models.resnet50(pretrained=True) 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_version=11 )常见问题及解决方案：
算子不支持：某些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 operator | ONNX算子不被TensorRT支持 | 使用插件或自定义层实现 |
| Invalid dimensions | 输入形状不匹配 | 检查ONNX模型输入输出形状 |
| Out of memory | 工作空间不足 | 增加max_workspace_size或减小batch size |
| FP16 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)), dtype=np.float32) output_host = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(output_bind_idx)), dtype=np.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_handle=stream.handle) cuda.memcpy_dtoh_async(output_host, output_device, stream) stream.synchronize()在实际项目中，从PyTorch到TensorRT的转换往往会遇到各种意想不到的问题。例如，某些自定义层可能无法直接转换，这时就需要编写TensorRT插件来实现相应功能。另外，不同版本的TensorRT对ONNX的支持程度也有所不同，遇到问题时查阅官方文档和社区讨论往往能快速找到解决方案。
