TensorRT入门指南：从模型优化到高效推理-CSDN博客

本文链接：https://blog.csdn.net/gitblog_00184/article/details/148393666

TensorRT入门指南：从模型优化到高效推理

TensorRT NVIDIA® TensorRT™ 是一个用于在 NVIDIA GPU 上进行高性能深度学习推理的软件开发工具包（SDK）。此代码库包含了 TensorRT 的开源组件项目地址: https://gitcode.com/gh_mirrors/tens/TensorRT

什么是TensorRT

TensorRT是NVIDIA推出的高性能深度学习推理SDK，它包含两个核心组件：

深度学习推理优化器：对训练好的深度学习模型进行优化
优化运行时：执行优化后的模型

TensorRT能够在保持模型精度的前提下，显著提升推理性能，降低延迟，减少内存占用。当你在PyTorch、TensorFlow等框架中完成模型训练后，TensorRT可以帮助你获得更高的推理吞吐量和更低的延迟。

TensorRT生态系统概览

TensorRT生态系统主要分为两大部分：

模型转换路径：将各种框架训练的模型转换为优化的TensorRT引擎
运行时部署：将优化后的TensorRT引擎部署到不同运行时环境

TensorRT生态系统

如何使用TensorRT

TensorRT支持多种工作流程，具体选择哪种取决于你的使用场景。抽象来看，从深度学习框架到TensorRT的部署流程如下：

TensorRT工作流程

在使用TensorRT时，你需要考虑以下五个关键问题：

模型保存格式：选择适合转换的模型格式
推理批大小：确定推理时的批处理规模
推理精度：选择合适的数值精度
转换路径：选择模型转换到TensorRT的方式
目标运行时：确定最终的部署环境

通过ONNX进行TensorRT转换的简单示例

我们将通过一个完整的示例，展示如何使用ONNX路径将预训练的ResNet50模型转换为优化的TensorRT引擎。

1. 准备模型格式

ONNX是TensorRT支持的通用模型格式之一。我们使用来自ONNX模型库的预训练ResNet50模型：

!wget https://download.onnxruntime.ai/onnx/models/resnet50.tar.gz -O resnet50.tar.gz
!tar xzf resnet50.tar.gz

2. 设置推理精度

TensorRT支持多种精度模式，包括TF32、FP32、FP16和INT8。这里我们使用FP32作为默认精度：

import numpy as np
PRECISION = np.float32
input_shape = (1, 3, 224, 224)  # 模型输入张量形状
dummy_input_batch = np.zeros(input_shape, dtype=PRECISION)

3. 使用ONNX转换路径

ONNX转换是TensorRT最通用且性能良好的自动转换路径之一。我们可以使用TensorRT自带的trtexec命令行工具进行转换：

!trtexec --onnx=resnet50/model.onnx --saveEngine=resnet_engine_intro.engine

这个命令会：

读取ONNX模型(resnet50/model.onnx)
转换为优化的TensorRT引擎(resnet_engine_intro.engine)
使用FP32精度(默认)

4. 运行推理

转换完成后，我们可以加载引擎并执行推理：

# 加载TensorRT引擎
with open("resnet_engine_intro.engine", "rb") as f:
    engine_data = f.read()
runtime = trt.Runtime(TRT_LOGGER)
engine = runtime.deserialize_cuda_engine(engine_data)

# 创建执行上下文
context = engine.create_execution_context()

# 分配输入/输出缓冲区
inputs, outputs, bindings = [], [], []
stream = cuda.Stream()

# 执行推理
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)

关键决策点详解

模型格式选择

TensorRT支持多种模型格式转换路径：

ONNX路径：通用性强，支持多框架
Torch-TRT路径：专为PyTorch设计，直接集成

精度选择策略

不同精度模式的特性比较：

| 精度 | 内存占用 | 计算速度 | 适用场景 | |------|---------|---------|---------| | FP32 | 高 | 中等 | 默认精度，兼容性好 | | FP16 | 中 | 快 | 大多数场景，性能与精度的平衡 | | INT8 | 低 | 最快 | 对延迟敏感，可接受轻微精度损失 |