ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

首页/资讯中心/资讯详情

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

发布时间：2026/9/25 17:04:12 尧图网页设计
从PyTorch到TensorRT：手把手教你将训练好的模型转成.engine文件（附完整代码）

从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 )

常见问题及解决方案：

  1. 算子不支持：某些PyTorch操作可能没有对应的ONNX实现。解决方法包括：

    • 使用更高版本的ONNX opset
    • 重写模型中使用的不支持算子
    • 添加自定义算子实现

  2. 动态维度处理：如果模型需要支持可变输入尺寸，务必在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 性能调优技巧

  1. 工作空间大小：适当增加max_workspace_size可以允许TensorRT尝试更多优化策略，但会占用更多内存。
  2. 层融合策略：通过config.set_tactic_sources控制TensorRT使用的优化策略。
  3. 时序缓存：保存优化策略缓存可以加速后续引擎构建。
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 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的支持程度也有所不同，遇到问题时查阅官方文档和社区讨论往往能快速找到解决方案。

网站建设 企业官网 郑州建站
返回列表