NVIDIA TensorRT 深度学习推理加速引擎详解

扫地的小何尚2025-03-26 8:21

NVIDIA TensorRT 深度学习推理加速引擎详解

文章目录

引言

在当今人工智能和深度学习快速发展的时代，模型推理性能已成为实际应用中的关键瓶颈。无论是需要实时响应的自动驾驶系统，还是要处理海量数据的推荐引擎，亦或是资源受限的边缘设备上的AI应用，都对推理速度、延迟和能耗提出了严苛的要求。NVIDIA TensorRT作为一个高性能的深度学习推理优化器和运行时环境，正是为解决这些挑战而生。

本文将全面深入地介绍NVIDIA TensorRT，从基本概念到安装部署，再到实际应用案例，帮助读者全面掌握这一强大工具。无论您是AI研究人员、深度学习工程师，还是对推理优化感兴趣的开发者，本文都将为您提供宝贵的指导和参考。

文章结构

本文将分为以下几个主要部分：

  1. TensorRT概述：介绍TensorRT的基本概念、核心功能和工作原理，帮助读者建立对TensorRT的整体认识。

  2. 安装与部署：详细说明TensorRT在不同平台（Linux、Windows、Docker等）上的安装步骤、系统要求和配置方法。

  3. 基础使用教程：通过实例讲解TensorRT的基本使用流程，包括模型转换、优化和部署等关键步骤。

  4. API详解与示例：分别介绍C++和Python API的使用方法，提供丰富的代码示例（附中文注释）。

  5. 应用场景实战：展示TensorRT在图像分类、目标检测、语义分割等实际应用中的使用方法和性能优势。

  6. 高级特性与优化：探讨TensorRT的高级功能，如量化、动态形状支持、自定义层等，以及性能优化的最佳实践。

  7. 常见问题与解决方案：总结使用TensorRT过程中可能遇到的问题及其解决方法。

让我们开始这段深入了解NVIDIA TensorRT的旅程，探索如何将您的深度学习模型转变为高性能的推理引擎。

第一部分：TensorRT概述

什么是TensorRT？

NVIDIA TensorRT是一个高性能深度学习推理优化器和运行时引擎，专为NVIDIA GPU设计，旨在提供低延迟和高吞吐量的模型推理。作为一个C++库，TensorRT可以读取通过各种框架（如TensorFlow、PyTorch、Caffe等）训练的模型，然后针对特定的GPU架构和应用需求进行优化，生成高度优化的推理引擎。

TensorRT不仅仅是一个简单的模型转换工具，它是一个完整的推理平台，提供了从模型优化到高效执行的全套功能。通过使用TensorRT，开发者可以显著提高深度学习应用的推理性能，同时减少延迟和能耗，这对于实时应用、边缘计算和大规模部署尤为重要。

TensorRT的核心功能和优势

1. 图优化

TensorRT通过一系列复杂的图优化技术来提高模型的执行效率：

  • 层融合（Layer Fusion）：将多个连续的层（如卷积+偏置+激活）合并为单一的优化操作，减少内存访问和计算开销。例如，将卷积层、批量归一化层和ReLU激活层融合为一个单一的优化层。

  • 内核自动调优（Kernel Auto-Tuning）：TensorRT会为特定的GPU架构和层配置选择最优的CUDA内核实现，确保最佳性能。

  • 动态张量内存管理：优化内存分配和重用，减少内存占用和数据传输开销。

  • 多精度优化：支持FP32（单精度浮点）、FP16（半精度浮点）和INT8（8位整数）等多种精度模式，可根据应用需求在精度和性能之间取得平衡。

2. 量化支持

TensorRT提供了强大的量化功能，可以将模型从FP32精度转换为FP16或INT8精度，在保持模型准确性的同时显著提高性能：

  • FP16精度：相比FP32，可以减少内存使用和带宽需求，同时在支持半精度计算的GPU上获得更高的计算吞吐量。

  • INT8量化：通过将浮点权重和激活值量化为8位整数，可以进一步提高性能，特别适合推理工作负载。TensorRT提供了校准工具，帮助确保量化后的模型保持准确性。

  • 稀疏性利用：能够利用模型中的结构化稀疏性来加速计算。

3. 动态形状支持

TensorRT支持动态输入形状，使得同一个优化引擎可以处理不同大小的输入：

  • 优化配置文件：允许为不同的输入形状范围定义优化配置文件。
  • 运行时形状调整：在推理时可以动态指定输入形状，无需重新构建引擎。

4. 多平台支持

TensorRT可以在各种NVIDIA平台上运行，从数据中心的高性能GPU到嵌入式设备：

  • 数据中心GPU：如Tesla V100、A100等。
  • 工作站GPU：如GeForce RTX系列、Quadro系列。
  • 嵌入式平台：如Jetson系列（Nano、TX2、Xavier、Orin等）。

5. 编程接口

TensorRT提供了多种编程接口，满足不同开发者的需求：

  • C++API：提供最完整的功能和最高的性能。
  • Python API：提供更便捷的开发体验，特别适合快速原型设计和实验。
  • 与深度学习框架的集成：如TensorFlow-TensorRT（TF-TRT）和ONNX解析器，简化了从训练框架到TensorRT的转换过程。

6. 性能优势

使用TensorRT可以获得显著的性能提升：

  • 低延迟：通过优化计算图和内存访问，减少推理延迟，对实时应用至关重要。
  • 高吞吐量：优化的批处理能力，适合需要处理大量数据的应用场景。
  • 能效提升：通过精度降低和计算优化，减少能耗，特别适合边缘设备和移动平台。

根据NVIDIA的测试，使用TensorRT优化后的模型相比原始框架可以获得数倍甚至数十倍的性能提升，具体提升幅度取决于模型结构、硬件平台和优化配置。

TensorRT的工作原理

TensorRT的工作流程可以分为以下几个关键阶段：

1. 模型导入

TensorRT首先需要导入训练好的深度学习模型。它支持多种导入方式：

  • 直接通过API构建网络：使用TensorRT的Network Definition API从头构建网络。
  • 解析预训练模型：使用内置的解析器（如ONNX解析器）导入预训练模型。
  • 通过框架特定的转换器：如TF-TRT（TensorFlow-TensorRT）。

2. 网络定义和优化

导入模型后，TensorRT会创建一个网络定义，然后应用一系列优化：

  • 层融合：识别可以合并的层模式，如卷积+批归一化+ReLU。
  • 张量精度校准：对于INT8量化，TensorRT会使用代表性数据集进行校准，确定最佳的量化参数。
  • 内核选择：为每个操作选择最优的CUDA内核实现。
  • 内存优化：规划最佳的内存分配和重用策略。

3. 引擎构建

优化后的网络被编译成一个高度优化的引擎，专门针对目标GPU架构：

  • 代码生成：生成针对特定GPU架构优化的CUDA代码。
  • 内存规划：确定张量的内存布局和分配策略。
  • 序列化：可以将构建好的引擎序列化到磁盘，以便后续快速加载。

4. 推理执行

使用构建好的引擎执行推理：

  • 内存分配：为输入和输出张量分配GPU内存。
  • 数据传输：将输入数据从主机内存复制到GPU内存。
  • 执行推理：调用优化后的CUDA内核执行计算。
  • 结果获取：将输出结果从GPU内存复制回主机内存。

5. 工作流程图示

以下是TensorRT工作流程的简化图示：

复制代码 
训练好的模型 → 模型导入 → 网络定义 → 优化 → 引擎构建 → 序列化 → 推理执行
                                  ↑
                            校准数据集（用于INT8量化）

TensorRT的应用场景

TensorRT广泛应用于各种需要高性能深度学习推理的场景：

1. 实时计算机视觉

  • 自动驾驶：需要低延迟处理来自多个传感器的数据。
  • 视频分析：实时监控和分析视频流。
  • 增强现实：需要在移动设备上实时处理视觉信息。

2. 自然语言处理

  • 语音识别：实时转录和理解语音输入。
  • 聊天机器人和虚拟助手：需要快速响应用户查询。
  • 机器翻译：实时翻译文本或语音。

3. 推荐系统

  • 电子商务：实时生成个性化产品推荐。
  • 内容平台：为用户推荐相关内容。

4. 边缘计算

  • 智能摄像头：在设备上进行对象检测和识别。
  • 机器人：实时感知和决策。
  • 物联网设备：在资源受限的设备上运行AI模型。

5. 医疗影像分析

  • 诊断辅助：快速分析医学图像。
  • 实时手术辅助：在手术过程中提供实时分析和指导。

TensorRT与其他推理框架的比较

与其他深度学习推理框架相比，TensorRT具有以下特点：

1. 与ONNX Runtime的比较

  • 性能：在NVIDIA GPU上，TensorRT通常比ONNX Runtime提供更高的性能。
  • 优化程度：TensorRT提供更深度的GPU特定优化。
  • 灵活性：ONNX Runtime支持更多硬件平台，而TensorRT专注于NVIDIA GPU。

2. 与TensorFlow Lite的比较

  • 目标平台：TensorFlow Lite主要针对移动和嵌入式设备，而TensorRT覆盖从嵌入式设备到高性能服务器的NVIDIA平台。
  • 性能：在NVIDIA硬件上，TensorRT通常提供更高的性能。
  • 生态系统：TensorFlow Lite与TensorFlow生态系统紧密集成。

3. 与PyTorch JIT/TorchScript的比较

  • 优化深度：TensorRT提供更深层次的硬件特定优化。
  • 使用便捷性：PyTorch JIT对PyTorch用户来说更为直接和简单。
  • 性能：在NVIDIA GPU上，TensorRT通常能提供更高的推理性能。

总结

NVIDIA TensorRT是一个功能强大的深度学习推理优化器和运行时引擎，通过一系列先进的优化技术，显著提高了深度学习模型的推理性能。它支持多种精度模式、动态形状和各种NVIDIA平台，为开发者提供了灵活而高效的推理解决方案。

无论是需要低延迟的实时应用，还是高吞吐量的批处理任务，或是在资源受限设备上的部署，TensorRT都能提供卓越的性能和效率。随着深度学习在各行各业的广泛应用，TensorRT作为推理优化的关键工具，将继续发挥重要作用。

在接下来的章节中，我们将详细介绍如何安装和配置TensorRT，以及如何使用它来优化和部署您的深度学习模型。

第二部分：TensorRT安装与部署指南

在开始使用NVIDIA TensorRT进行深度学习模型优化和推理之前，我们首先需要正确安装和配置TensorRT环境。本章将详细介绍TensorRT的系统要求、不同平台的安装方法以及相关配置步骤，帮助您快速搭建TensorRT开发环境。

系统要求和前提条件

在安装TensorRT之前，请确保您的系统满足以下基本要求：

硬件要求

  • 支持的GPU：NVIDIA GPU，包括：

    • 数据中心GPU：Tesla、A100、A30、A10等
    • 工作站GPU：GeForce RTX系列、Quadro系列
    • 嵌入式平台：Jetson系列（Nano、TX2、Xavier、Orin等）

  • 最低GPU计算能力：

    • 对于TensorRT 8.x版本：支持CUDA计算能力3.0或更高版本的GPU
    • 对于最新特性（如稀疏性、多实例GPU等）：建议使用计算能力7.0或更高版本的GPU

软件要求

  • 操作系统：

    • Linux：Ubuntu 18.04/20.04/22.04 LTS、RHEL/CentOS 7/8
    • Windows：Windows 10/11、Windows Server 2019/2022
    • 对于Jetson平台：JetPack SDK（包含特定版本的L4T操作系统）

  • CUDA工具包：

    • TensorRT 8.x需要CUDA 11.0或更高版本
    • 建议安装最新的CUDA版本以获得最佳性能

  • cuDNN库：

    • 与所安装的CUDA版本兼容的cuDNN版本
    • 通常需要cuDNN 8.x或更高版本

  • Python（可选，用于Python API）：

    • Python 3.6或更高版本
    • 推荐使用Python 3.8或3.9以获得最佳兼容性

不同平台的安装方法

Linux平台安装

Linux是TensorRT的主要支持平台，提供了多种安装方式。

方法1：使用包管理器安装（推荐）

通过NVIDIA软件包存储库安装是最简单的方法：

bash 复制代码 
# 添加NVIDIA软件包存储库
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"

# 更新软件包列表
sudo apt-get update

# 安装TensorRT
sudo apt-get install libnvinfer8 libnvinfer-plugin8 libnvparsers8 libnvonnxparsers8 libnvinfer-dev libnvinfer-plugin-dev libnvparsers-dev libnvonnxparsers-dev libnvinfer-samples tensorrt-dev

# 安装Python API（可选）
sudo apt-get install python3-libnvinfer
方法2：使用TAR包安装

如果您需要特定版本或无法使用包管理器，可以下载TAR包手动安装：

  1. 访问NVIDIA TensorRT下载页面
  2. 下载适合您系统的TAR包
  3. 解压TAR包：
bash 复制代码 
tar -xzvf TensorRT-8.x.x.x.Linux.x86_64-gnu.cuda-xx.x.tar.gz
  1. 设置环境变量：
bash 复制代码 
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:<TensorRT-安装路径>/lib
  1. 安装Python包（可选）：
bash 复制代码 
cd <TensorRT-安装路径>/python
pip install tensorrt-*-cp3x-none-linux_x86_64.whl

cd <TensorRT-安装路径>/graphsurgeon
pip install graphsurgeon-0.4.x-py2.py3-none-any.whl

cd <TensorRT-安装路径>/onnx_graphsurgeon
pip install onnx_graphsurgeon-0.3.x-py2.py3-none-any.whl
方法3：使用NGC容器（适用于容器化部署）

NVIDIA NGC提供了预装TensorRT的Docker容器：

bash 复制代码 
# 拉取TensorRT容器
docker pull nvcr.io/nvidia/tensorrt:22.12-py3

# 运行容器
docker run --gpus all -it --rm nvcr.io/nvidia/tensorrt:22.12-py3

Windows平台安装

Windows平台上安装TensorRT的步骤如下：

  1. 确保已安装兼容版本的CUDA和cuDNN

  2. 访问NVIDIA TensorRT下载页面

  3. 下载适用于Windows的ZIP包

  4. 解压ZIP包到所需位置

  5. 将TensorRT的bin目录添加到系统PATH环境变量：

    C:\Program Files\NVIDIA\TensorRT\lib

  6. 安装Python包（可选）：

cmd 复制代码 
cd C:\Program Files\NVIDIA\TensorRT\python
pip install tensorrt-*-cp3x-none-win_amd64.whl

cd C:\Program Files\NVIDIA\TensorRT\graphsurgeon
pip install graphsurgeon-0.4.x-py2.py3-none-any.whl

cd C:\Program Files\NVIDIA\TensorRT\onnx_graphsurgeon
pip install onnx_graphsurgeon-0.3.x-py2.py3-none-any.whl

使用Docker安装

使用Docker是跨平台使用TensorRT的便捷方式，避免了环境配置的复杂性：

bash 复制代码 
# 拉取官方TensorRT容器
docker pull nvcr.io/nvidia/tensorrt:22.12-py3

# 运行容器，挂载本地目录以便访问模型和数据
docker run --gpus all -it --rm -v /path/to/your/models:/models nvcr.io/nvidia/tensorrt:22.12-py3

Jetson平台安装

对于NVIDIA Jetson设备（如Jetson Nano、TX2、Xavier或Orin），TensorRT已包含在JetPack SDK中：

  1. 下载并安装最新版本的JetPack SDK
  2. 使用NVIDIA SDK Manager安装JetPack到您的Jetson设备
  3. TensorRT将自动安装为JetPack的一部分

验证安装

安装完成后，可以通过以下方法验证TensorRT是否正确安装：

检查TensorRT版本

bash 复制代码 
# 在Linux终端或Windows命令提示符中
dpkg -l | grep TensorRT  # 仅适用于Debian/Ubuntu

# 或者使用Python验证
python3 -c "import tensorrt as trt; print(trt.__version__)"

运行示例程序

TensorRT安装包中包含多个示例程序，可以用来验证安装：

bash 复制代码 
# 在Linux上
cd <TensorRT-安装路径>/samples
make
cd bin
./sample_mnist

# 在Windows上
cd C:\Program Files\NVIDIA\TensorRT\samples
msbuild samples_vs2019.sln
cd C:\Program Files\NVIDIA\TensorRT\bin
sample_mnist.exe

常见安装问题及解决方案

1. 库依赖问题

问题：安装后出现库依赖错误。

解决方案：

bash 复制代码 
# 检查并安装缺失的依赖
sudo apt-get install libcudnn8 libcublas10

2. Python包导入错误

问题：无法导入tensorrt模块。

解决方案：

bash 复制代码 
# 确认Python包安装路径
pip list | grep tensorrt

# 重新安装Python包
pip uninstall tensorrt
cd <TensorRT-安装路径>/python
pip install tensorrt-*-cp3x-none-linux_x86_64.whl

3. CUDA版本不兼容

问题：TensorRT与已安装的CUDA版本不兼容。

解决方案：

bash 复制代码 
# 检查CUDA版本
nvcc --version

# 安装兼容的TensorRT版本或更新CUDA

升级TensorRT

当需要升级到新版本的TensorRT时，请按照以下步骤操作：

Linux平台升级

bash 复制代码 
# 使用包管理器升级
sudo apt-get update
sudo apt-get upgrade libnvinfer* tensorrt*

# 或者使用TAR包手动升级
# 1. 下载新版本的TAR包
# 2. 解压到新目录
# 3. 更新环境变量指向新目录

Windows平台升级

  1. 下载新版本的ZIP包
  2. 解压到新位置或覆盖旧版本
  3. 更新PATH环境变量（如果安装位置变更）
  4. 重新安装Python包

卸载TensorRT

如果需要卸载TensorRT，可以按照以下步骤操作：

Linux平台卸载

bash 复制代码 
# 使用包管理器卸载
sudo apt-get remove --purge libnvinfer* tensorrt*

# 如果是手动安装，删除安装目录并移除环境变量设置
rm -rf <TensorRT-安装路径>
# 从~/.bashrc或~/.profile中移除相关环境变量设置

Windows平台卸载

  1. 从控制面板的"程序和功能"中卸载TensorRT
  2. 或者直接删除安装目录
  3. 从系统PATH环境变量中移除TensorRT路径

总结

本章详细介绍了TensorRT的安装和部署方法，涵盖了不同平台（Linux、Windows、Docker和Jetson）的安装步骤、验证方法以及常见问题的解决方案。通过正确安装和配置TensorRT，您已经为后续的深度学习模型优化和部署工作奠定了基础。

在下一章中，我们将开始探索TensorRT的基本使用流程，包括如何导入模型、优化网络以及执行推理。

第三部分：TensorRT基础使用教程

在完成TensorRT的安装和配置后，本章将带您了解TensorRT的基本使用流程。我们将通过详细的步骤和丰富的代码示例，展示如何使用TensorRT来优化深度学习模型并执行高效推理。

TensorRT工作流程概述

使用TensorRT进行模型优化和推理的典型工作流程包括以下几个主要步骤：

  1. 模型准备：获取预训练模型或从深度学习框架导出模型
  2. 创建TensorRT网络：通过解析器导入模型或使用API构建网络
  3. 构建优化引擎：配置优化参数并构建TensorRT引擎
  4. 序列化引擎：将引擎保存到磁盘以便后续使用
  5. 执行推理：加载引擎并执行高效推理

下面我们将详细介绍每个步骤的具体操作方法。

模型准备

在使用TensorRT之前，您需要准备好要优化的深度学习模型。TensorRT支持多种模型导入方式，最常用的是通过ONNX格式。

从PyTorch导出ONNX模型

python 复制代码 
# 从PyTorch导出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",     # 输出文件名
                  export_params=True,  # 存储训练好的参数权重
                  opset_version=13,    # ONNX算子集版本
                  do_constant_folding=True,  # 是否执行常量折叠优化
                  input_names=["input"],     # 输入名称
                  output_names=["output"],   # 输出名称
                  dynamic_axes={"input": {0: "batch_size"},    # 动态轴
                                "output": {0: "batch_size"}})

print("ONNX模型已导出到resnet50.onnx")

从TensorFlow导出ONNX模型

对于TensorFlow模型，可以使用tf2onnx工具进行转换：

python 复制代码 
# 从TensorFlow导出ONNX模型
import tensorflow as tf
import tf2onnx

# 加载预训练模型
model = tf.keras.applications.ResNet50(weights='imagenet')

# 转换为ONNX
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
output_path = "resnet50.onnx"

model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13, output_path=output_path)
print("ONNX模型已导出到resnet50.onnx")

使用C++ API创建TensorRT引擎

下面是使用C++ API从ONNX模型创建TensorRT引擎的完整示例：

cpp 复制代码 
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include <fstream>
#include <iostream>
#include <memory>

// 自定义日志记录器类
class Logger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        // 忽略INFO级别的消息
        if (severity <= Severity::kWARNING)
            std::cout << msg << std::endl;
    }
} logger;

// 智能指针辅助函数，用于自动释放TensorRT对象
template <typename T>
struct TRTDestroy {
    void operator()(T* obj) const {
        if (obj)
            obj->destroy();
    }
};
template <typename T>
using TRTUniquePtr = std::unique_ptr<T, TRTDestroy<T>>;

// 从文件读取二进制数据
std::vector<unsigned char> readFile(const std::string& fileName) {
    std::ifstream file(fileName, std::ios::binary | std::ios::ate);
    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<unsigned char> buffer(size);
    file.read(reinterpret_cast<char*>(buffer.data()), size);
    file.close();
    return buffer;
}

// 将引擎序列化并保存到文件
void saveEngine(const nvinfer1::ICudaEngine& engine, const std::string& fileName) {
    TRTUniquePtr<nvinfer1::IHostMemory> serializedEngine{engine.serialize()};
    std::ofstream file(fileName, std::ios::binary);
    file.write(reinterpret_cast<const char*>(serializedEngine->data()), serializedEngine->size());
    file.close();
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        std::cerr << "用法: " << argv[0] << " <onnx_model> <engine_file>" << std::endl;
        return 1;
    }
    
    std::string onnxFile = argv[1];
    std::string engineFile = argv[2];
    
    // 创建构建器
    TRTUniquePtr<nvinfer1::IBuilder> builder{nvinfer1::createInferBuilder(logger)};
    if (!builder) {
        std::cerr << "创建构建器失败" << std::endl;
        return 1;
    }
    
    // 创建网络定义
    const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    TRTUniquePtr<nvinfer1::INetworkDefinition> network{builder->createNetworkV2(explicitBatch)};
    if (!network) {
        std::cerr << "创建网络定义失败" << std::endl;
        return 1;
    }
    
    // 创建ONNX解析器
    TRTUniquePtr<nvonnxparser::IParser> parser{nvonnxparser::createParser(*network, logger)};
    if (!parser) {
        std::cerr << "创建ONNX解析器失败" << std::endl;
        return 1;
    }
    
    // 解析ONNX模型
    auto onnxData = readFile(onnxFile);
    if (!parser->parse(onnxData.data(), onnxData.size())) {
        std::cerr << "解析ONNX模型失败" << std::endl;
        return 1;
    }
    
    // 创建构建配置
    TRTUniquePtr<nvinfer1::IBuilderConfig> config{builder->createBuilderConfig()};
    if (!config) {
        std::cerr << "创建构建配置失败" << std::endl;
        return 1;
    }
    
    // 设置最大工作空间大小（1GB）
    config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1ULL << 30);
    
    // 启用FP16精度（如果GPU支持）
    if (builder->platformHasFastFp16()) {
        std::cout << "启用FP16精度" << std::endl;
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
    }
    
    // 构建引擎
    std::cout << "构建TensorRT引擎..." << std::endl;
    TRTUniquePtr<nvinfer1::ICudaEngine> engine{builder->buildEngineWithConfig(*network, *config)};
    if (!engine) {
        std::cerr << "构建引擎失败" << std::endl;
        return 1;
    }
    
    // 保存引擎到文件
    std::cout << "将引擎序列化到: " << engineFile << std::endl;
    saveEngine(*engine, engineFile);
    
    std::cout << "引擎创建成功!" << std::endl;
    return 0;
}

编译上述代码：

bash 复制代码 
g++ -o build_engine build_engine.cpp -I/usr/local/cuda/include -I/usr/include/x86_64-linux-gnu -L/usr/local/cuda/lib64 -lnvinfer -lnvonnxparser -lcudart

使用Python API创建TensorRT引擎

以下是使用Python API从ONNX模型创建TensorRT引擎的示例：

python 复制代码 
import os
import tensorrt as trt
import numpy as np

def build_engine_from_onnx(onnx_file_path, engine_file_path, precision='fp32'):
    """
    从ONNX模型构建TensorRT引擎
    
    参数:
        onnx_file_path: ONNX模型文件路径
        engine_file_path: 输出的TensorRT引擎文件路径
        precision: 精度模式，可选 'fp32', 'fp16', 'int8'
    
    返回:
        构建好的TensorRT引擎
    """
    # 创建logger和builder
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(TRT_LOGGER)
    
    # 创建网络定义
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    
    # 创建ONNX解析器
    parser = trt.OnnxParser(network, TRT_LOGGER)
    
    # 解析ONNX文件
    with open(onnx_file_path, 'rb') as model:
        if not parser.parse(model.read()):
            print('ERROR: 解析ONNX模型失败')
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None
    
    # 创建构建配置
    config = builder.create_builder_config()
    
    # 设置最大工作空间大小（1GB）
    config.max_workspace_size = 1 << 30
    
    # 根据指定精度设置标志
    if precision == 'fp16' and builder.platform_has_fast_fp16:
        print('启用FP16精度')
        config.set_flag(trt.BuilderFlag.FP16)
    elif precision == 'int8' and builder.platform_has_fast_int8:
        print('启用INT8精度')
        config.set_flag(trt.BuilderFlag.INT8)
        # 注意：INT8需要校准器，这里简化处理
    
    # 构建引擎
    print('构建TensorRT引擎...')
    serialized_engine = builder.build_serialized_network(network, config)
    
    # 保存引擎到文件
    with open(engine_file_path, 'wb') as f:
        f.write(serialized_engine)
    
    # 创建运行时并反序列化引擎
    runtime = trt.Runtime(TRT_LOGGER)
    engine = runtime.deserialize_cuda_engine(serialized_engine)
    
    print(f'引擎已保存到: {engine_file_path}')
    return engine

# 使用示例
if __name__ == '__main__':
    onnx_model = 'resnet50.onnx'
    engine_file = 'resnet50.engine'
    
    # 构建FP16精度的引擎
    engine = build_engine_from_onnx(onnx_model, engine_file, precision='fp16')

使用TensorRT执行推理

C++推理示例

以下是使用C++ API加载TensorRT引擎并执行推理的示例：

cpp 复制代码 
#include "NvInfer.h"
#include <fstream>
#include <iostream>
#include <memory>
#include <vector>
#include <cuda_runtime_api.h>

// 自定义日志记录器类
class Logger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        if (severity <= Severity::kWARNING)
            std::cout << msg << std::endl;
    }
} logger;

// 智能指针辅助函数
template <typename T>
struct TRTDestroy {
    void operator()(T* obj) const {
        if (obj)
            obj->destroy();
    }
};
template <typename T>
using TRTUniquePtr = std::unique_ptr<T, TRTDestroy<T>>;

// 从文件读取引擎
std::vector<unsigned char> readEngineFile(const std::string& fileName) {
    std::ifstream file(fileName, std::ios::binary | std::ios::ate);
    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<unsigned char> buffer(size);
    file.read(reinterpret_cast<char*>(buffer.data()), size);
    file.close();
    return buffer;
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cerr << "用法: " << argv[0] << " <engine_file>" << std::endl;
        return 1;
    }
    
    std::string engineFile = argv[1];
    
    // 创建运行时
    TRTUniquePtr<nvinfer1::IRuntime> runtime{nvinfer1::createInferRuntime(logger)};
    if (!runtime) {
        std::cerr << "创建运行时失败" << std::endl;
        return 1;
    }
    
    // 加载引擎
    auto engineData = readEngineFile(engineFile);
    TRTUniquePtr<nvinfer1::ICudaEngine> engine{
        runtime->deserializeCudaEngine(engineData.data(), engineData.size())
    };
    if (!engine) {
        std::cerr << "反序列化引擎失败" << std::endl;
        return 1;
    }
    
    // 创建执行上下文
    TRTUniquePtr<nvinfer1::IExecutionContext> context{engine->createExecutionContext()};
    if (!context) {
        std::cerr << "创建执行上下文失败" << std::endl;
        return 1;
    }
    
    // 获取绑定索引
    int inputIndex = engine->getBindingIndex("input");
    int outputIndex = engine->getBindingIndex("output");
    
    // 获取输入和输出维度
    auto inputDims = engine->getBindingDimensions(inputIndex);
    auto outputDims = engine->getBindingDimensions(outputIndex);
    
    // 计算输入和输出大小
    int batchSize = 1;  // 假设批处理大小为1
    int inputSize = batchSize;
    int outputSize = batchSize;
    
    for (int i = 1; i < inputDims.nbDims; i++)
        inputSize *= inputDims.d[i];
    
    for (int i = 1; i < outputDims.nbDims; i++)
        outputSize *= outputDims.d[i];
    
    // 分配主机和设备内存
    std::vector<float> inputData(inputSize, 0.5f);  // 示例输入数据
    std::vector<float> outputData(outputSize);
    
    void* deviceBuffers[2];
    cudaMalloc(&deviceBuffers[inputIndex], inputSize * sizeof(float));
    cudaMalloc(&deviceBuffers[outputIndex], outputSize * sizeof(float));
    
    // 创建CUDA流
    cudaStream_t stream;
    cudaStreamCreate(&stream);
    
    // 将输入数据复制到设备
    cudaMemcpyAsync(deviceBuffers[inputIndex], inputData.data(), 
                   inputSize * sizeof(float), cudaMemcpyHostToDevice, stream);
    
    // 执行推理
    context->enqueueV2(deviceBuffers, stream, nullptr);
    
    // 将输出数据复制回主机
    cudaMemcpyAsync(outputData.data(), deviceBuffers[outputIndex],
                   outputSize * sizeof(float), cudaMemcpyDeviceToHost, stream);
    
    // 同步流
    cudaStreamSynchronize(stream);
    
    // 输出结果（仅显示前5个值）
    std::cout << "推理结果 (前5个值):" << std::endl;
    for (int i = 0; i < std::min(5, outputSize); i++) {
        std::cout << "  " << i << ": " << outputData[i] << std::endl;
    }
    
    // 清理资源
    cudaFree(deviceBuffers[inputIndex]);
    cudaFree(deviceBuffers[outputIndex]);
    cudaStreamDestroy(stream);
    
    return 0;
}

Python推理示例

以下是使用Python API加载TensorRT引擎并执行推理的示例：

python 复制代码 
import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit

class TensorRTInference:
    def __init__(self, engine_path):
        """
        初始化TensorRT推理器
        
        参数:
            engine_path: TensorRT引擎文件路径
        """
        # 创建logger和运行时
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        
        # 加载引擎
        with open(engine_path, 'rb') as f:
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        
        # 创建执行上下文
        self.context = self.engine.create_execution_context()
        
        # 创建CUDA流
        self.stream = cuda.Stream()
        
        # 获取绑定索引
        self.input_binding_idxs = []
        self.output_binding_idxs = []
        
        for i in range(self.engine.num_bindings):
            if self.engine.binding_is_input(i):
                self.input_binding_idxs.append(i)
            else:
                self.output_binding_idxs.append(i)
        
        # 为输出分配设备内存
        self.output_buffers = {}
        for i in self.output_binding_idxs:
            name = self.engine.get_binding_name(i)
            dtype = trt.nptype(self.engine.get_binding_dtype(i))
            shape = self.context.get_binding_shape(i)
            size = trt.volume(shape)
            
            # 分配主机和设备内存
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            
            self.output_buffers[name] = {
                'host': host_mem,
                'device': device_mem,
                'shape': shape
            }
    
    def infer(self, input_data):
        """
        执行推理
        
        参数:
            input_data: 输入数据字典，键为输入名称，值为numpy数组
            
        返回:
            输出数据字典，键为输出名称，值为numpy数组
        """
        # 准备输入绑定
        bindings = []
        input_buffers = {}
        
        # 添加输入绑定
        for i in self.input_binding_idxs:
            name = self.engine.get_binding_name(i)
            
            if name not in input_data:
                raise ValueError(f"输入 '{name}' 未提供")
            
            data = input_data[name]
            
            # 设置动态形状（如果需要）
            if self.engine.is_shape_binding(i) or -1 in self.context.get_binding_shape(i):
                self.context.set_binding_shape(i, data.shape)
            
            # 分配主机和设备内存
            dtype = trt.nptype(self.engine.get_binding_dtype(i))
            host_mem = cuda.pagelocked_empty(data.size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            
            # 复制数据
            np.copyto(host_mem, data.ravel())
            cuda.memcpy_htod_async(device_mem, host_mem, self.stream)
            
            input_buffers[name] = {
                'host': host_mem,
                'device': device_mem
            }
            
            bindings.append(int(device_mem))
        
        # 添加输出绑定
        for i in self.output_binding_idxs:
            name = self.engine.get_binding_name(i)
            bindings.append(int(self.output_buffers[name]['device']))
        
        # 执行推理
        self.context.execute_async_v2(bindings=bindings, stream_handle=self.stream.handle)
        
        # 将输出从设备复制到主机
        outputs = {}
        for i in self.output_binding_idxs:
            name = self.engine.get_binding_name(i)
            output_shape = self.context.get_binding_shape(i)
            output_size = trt.volume(output_shape)
            
            cuda.memcpy_dtoh_async(self.output_buffers[name]['host'], 
                                  self.output_buffers[name]['device'], 
                                  self.stream)
            
            # 重塑输出
            outputs[name] = self.output_buffers[name]['host'].reshape(output_shape)
        
        # 同步流
        self.stream.synchronize()
        
        # 清理输入缓冲区
        for name, buffer in input_buffers.items():
            buffer['device'].free()
        
        return outputs
    
    def __del__(self):
        """清理资源"""
        for name, buffer in self.output_buffers.items():
            buffer['device'].free()

# 使用示例
if __name__ == '__main__':
    # 加载引擎
    engine_path = 'resnet50.engine'
    inference = TensorRTInference(engine_path)
    
    # 准备输入数据（示例：随机数据）
    input_shape = (1, 3, 224, 224)  # 批大小=1, 3通道, 224x224图像
    input_data = {
        'input': np.random.random(input_shape).astype(np.float32)
    }
    
    # 执行推理
    outputs = inference.infer(input_data)
    
    # 处理输出（假设输出名为'output'）
    if 'output' in outputs:
        # 获取前5个最高概率的类别
        probs = outputs['output'][0]
        top5_indices = np.argsort(probs)[-5:][::-1]
        
        print("Top 5预测结果:")
        for i, idx in enumerate(top5_indices):
            print(f"  {i+1}. 类别 {idx}: {probs[idx]:.6f}")

高级配置选项

启用FP16精度

TensorRT支持FP16（半精度浮点）计算，可以显著提高性能，同时只有轻微的精度损失：

python 复制代码 
# 在Python中启用FP16
config = builder.create_builder_config()
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)
cpp 复制代码 
// 在C++中启用FP16
if (builder->platformHasFastFp16()) {
    config->setFlag(nvinfer1::BuilderFlag::kFP16);
}

启用INT8量化

INT8量化可以进一步提高性能，但需要校准数据来确保精度：

python 复制代码 
# 在Python中启用INT8（简化版，实际使用需要校准器）
config = builder.create_builder_config()
if builder.platform_has_fast_int8:
    config.set_flag(trt.BuilderFlag.INT8)
    # 这里需要添加校准器配置
cpp 复制代码 
// 在C++中启用INT8（简化版，实际使用需要校准器）
if (builder->platformHasFastInt8()) {
    config->setFlag(nvinfer1::BuilderFlag::kINT8);
    // 这里需要添加校准器配置
}

配置动态形状

TensorRT支持动态输入形状，使同一引擎可以处理不同大小的输入：

python 复制代码 
# 在Python中配置动态形状
profile = builder.create_optimization_profile()
profile.set_shape(
    "input",                          # 输入名称
    (1, 3, 224, 224),                 # 最小形状
    (4, 3, 224, 224),                 # 优化形状
    (8, 3, 224, 224)                  # 最大形状
)
config.add_optimization_profile(profile)
cpp 复制代码 
// 在C++中配置动态形状
auto profile = builder->createOptimizationProfile();
profile->setDimensions(
    "input",                          // 输入名称
    nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, 3, 224, 224));
profile->setDimensions(
    "input",                          // 输入名称
    nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(4, 3, 224, 224));
profile->setDimensions(
    "input",                          // 输入名称
    nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(8, 3, 224, 224));
config->addOptimizationProfile(profile);

性能优化技巧

1. 批处理推理

对多个输入进行批处理可以提高吞吐量：

python 复制代码 
# 批处理推理示例
batch_size = 8
input_data = np.random.random((batch_size, 3, 224, 224)).astype(np.float32)

2. 使用CUDA图

对于固定输入大小的推理，可以使用CUDA图进一步提高性能：

python 复制代码 
# 使用CUDA图的示例（需要CUDA 10.0或更高版本）
import pycuda.driver as cuda
# 创建CUDA图
stream = cuda.Stream()
cuda.graph_begin_capture(stream, cuda.graph_capture_mode.GLOBAL)
# 执行推理
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
# 结束捕获
graph = cuda.graph_end_capture(stream)
# 实例化图
graph_exec = cuda.graph_instantiate(graph)
# 使用图执行推理
graph_exec.launch(stream)
stream.synchronize()

3. 内存优化

减少主机和设备之间的数据传输：

python 复制代码 
# 使用页锁定内存
host_input = cuda.pagelocked_empty(input_size, dtype=np.float32)
host_output = cuda.pagelocked_empty(output_size, dtype=np.float32)

4. 选择最佳精度

根据应用需求选择合适的精度：

  • FP32：最高精度，但性能较低
  • FP16：平衡精度和性能
  • INT8：最高性能，但精度较低，需要校准

总结

本章详细介绍了TensorRT的基本使用流程，包括模型准备、引擎构建和推理执行。我们提供了完整的C++和Python代码示例，展示了如何从ONNX模型创建TensorRT引擎，以及如何使用引擎执行高效推理。此外，我们还介绍了一些高级配置选项和性能优化技巧，帮助您充分发挥TensorRT的性能潜力。

在下一章中，我们将深入探讨TensorRT的API详解和更多实际应用场景，帮助您更全面地掌握TensorRT的使用方法。

第四部分：TensorRT API详解与示例

在前面的章节中，我们已经了解了TensorRT的基本概念、安装方法和基础使用流程。本章将深入探讨TensorRT的API，包括C++和Python两种接口，并提供详细的代码示例和中文注释，帮助读者全面掌握TensorRT的编程方法。

TensorRT API概述

TensorRT提供了两套主要的API：

  1. C++ API：提供最完整的功能和最高的性能，适合对性能要求极高或需要深度定制的应用场景。
  2. Python API：提供更便捷的开发体验，特别适合快速原型设计和实验，同时保持与C++ API相似的结构。

两种API都遵循相似的工作流程：创建网络定义、构建优化引擎、执行推理。下面我们将分别详细介绍这两套API的使用方法。

C++ API详解

核心类和对象

TensorRT C++ API的核心类包括：

  • ILogger：日志接口，用于记录TensorRT操作过程中的消息
  • IBuilder：构建器，用于创建网络定义和构建引擎
  • INetworkDefinition：网络定义，表示深度学习模型的计算图
  • IBuilderConfig：构建配置，控制引擎构建过程中的优化选项
  • ICudaEngine：CUDA引擎，表示优化后的模型
  • IExecutionContext：执行上下文，用于执行推理

详细的C++ API示例

下面是一个完整的C++ API使用示例，包含详细的中文注释：

cpp 复制代码 
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include <fstream>
#include <iostream>
#include <memory>
#include <vector>
#include <cuda_runtime_api.h>

// 自定义日志记录器类，继承自ILogger接口
class Logger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        // 根据日志级别过滤消息，只显示警告及以上级别
        if (severity <= Severity::kWARNING)
            std::cout << msg << std::endl;
    }
} logger;

// 智能指针辅助函数，用于自动释放TensorRT对象
template <typename T>
struct TRTDestroy {
    void operator()(T* obj) const {
        if (obj)
            obj->destroy();
    }
};
template <typename T>
using TRTUniquePtr = std::unique_ptr<T, TRTDestroy<T>>;

// 从文件读取二进制数据的辅助函数
std::vector<unsigned char> readFile(const std::string& fileName) {
    std::ifstream file(fileName, std::ios::binary | std::ios::ate);
    if (!file) {
        std::cerr << "无法打开文件: " << fileName << std::endl;
        return {};
    }
    
    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<unsigned char> buffer(size);
    file.read(reinterpret_cast<char*>(buffer.data()), size);
    file.close();
    return buffer;
}

// 将引擎序列化并保存到文件的辅助函数
void saveEngine(const nvinfer1::ICudaEngine& engine, const std::string& fileName) {
    // 序列化引擎
    TRTUniquePtr<nvinfer1::IHostMemory> serializedEngine{engine.serialize()};
    
    // 将序列化数据写入文件
    std::ofstream file(fileName, std::ios::binary);
    file.write(reinterpret_cast<const char*>(serializedEngine->data()), serializedEngine->size());
    file.close();
}

// 主函数：从ONNX模型构建TensorRT引擎并执行推理
int main(int argc, char* argv[]) {
    if (argc < 3) {
        std::cerr << "用法: " << argv[0] << " <onnx_model> <engine_file>" << std::endl;
        return 1;
    }
    
    std::string onnxFile = argv[1];
    std::string engineFile = argv[2];
    
    // 第一部分：构建TensorRT引擎
    
    // 创建构建器
    TRTUniquePtr<nvinfer1::IBuilder> builder{nvinfer1::createInferBuilder(logger)};
    if (!builder) {
        std::cerr << "创建构建器失败" << std::endl;
        return 1;
    }
    
    // 创建网络定义，使用显式批处理维度
    const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    TRTUniquePtr<nvinfer1::INetworkDefinition> network{builder->createNetworkV2(explicitBatch)};
    if (!network) {
        std::cerr << "创建网络定义失败" << std::endl;
        return 1;
    }
    
    // 创建ONNX解析器
    TRTUniquePtr<nvonnxparser::IParser> parser{nvonnxparser::createParser(*network, logger)};
    if (!parser) {
        std::cerr << "创建ONNX解析器失败" << std::endl;
        return 1;
    }
    
    // 解析ONNX模型
    auto onnxData = readFile(onnxFile);
    if (onnxData.empty()) {
        std::cerr << "读取ONNX模型失败" << std::endl;
        return 1;
    }
    
    if (!parser->parse(onnxData.data(), onnxData.size())) {
        std::cerr << "解析ONNX模型失败" << std::endl;
        for (int i = 0; i < parser->getNbErrors(); ++i) {
            std::cerr << parser->getError(i)->desc() << std::endl;
        }
        return 1;
    }
    
    // 创建构建配置
    TRTUniquePtr<nvinfer1::IBuilderConfig> config{builder->createBuilderConfig()};
    if (!config) {
        std::cerr << "创建构建配置失败" << std::endl;
        return 1;
    }
    
    // 设置最大工作空间大小（1GB）
    config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1ULL << 30);
    
    // 启用FP16精度（如果GPU支持）
    if (builder->platformHasFastFp16()) {
        std::cout << "启用FP16精度" << std::endl;
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
    }
    
    // 构建引擎
    std::cout << "构建TensorRT引擎..." << std::endl;
    TRTUniquePtr<nvinfer1::ICudaEngine> engine{builder->buildEngineWithConfig(*network, *config)};
    if (!engine) {
        std::cerr << "构建引擎失败" << std::endl;
        return 1;
    }
    
    // 保存引擎到文件
    std::cout << "将引擎序列化到: " << engineFile << std::endl;
    saveEngine(*engine, engineFile);
    
    // 第二部分：使用TensorRT引擎执行推理
    
    // 创建执行上下文
    TRTUniquePtr<nvinfer1::IExecutionContext> context{engine->createExecutionContext()};
    if (!context) {
        std::cerr << "创建执行上下文失败" << std::endl;
        return 1;
    }
    
    // 获取输入和输出的绑定索引
    int inputIndex = engine->getBindingIndex("input");  // 假设输入名称为"input"
    int outputIndex = engine->getBindingIndex("output");  // 假设输出名称为"output"
    
    if (inputIndex == -1 || outputIndex == -1) {
        std::cerr << "无法找到输入或输出绑定" << std::endl;
        return 1;
    }
    
    // 获取输入和输出维度
    auto inputDims = engine->getBindingDimensions(inputIndex);
    auto outputDims = engine->getBindingDimensions(outputIndex);
    
    // 计算输入和输出大小
    int batchSize = 1;  // 假设批处理大小为1
    int inputSize = batchSize;
    int outputSize = batchSize;
    
    for (int i = 1; i < inputDims.nbDims; i++)
        inputSize *= inputDims.d[i];
    
    for (int i = 1; i < outputDims.nbDims; i++)
        outputSize *= outputDims.d[i];
    
    // 分配主机内存
    std::vector<float> inputData(inputSize, 0.5f);  // 示例输入数据，全部填充为0.5
    std::vector<float> outputData(outputSize);
    
    // 分配设备内存
    void* deviceBuffers[2];
    cudaMalloc(&deviceBuffers[inputIndex], inputSize * sizeof(float));
    cudaMalloc(&deviceBuffers[outputIndex], outputSize * sizeof(float));
    
    // 创建CUDA流
    cudaStream_t stream;
    cudaStreamCreate(&stream);
    
    // 将输入数据从主机复制到设备
    cudaMemcpyAsync(deviceBuffers[inputIndex], inputData.data(), 
                   inputSize * sizeof(float), cudaMemcpyHostToDevice, stream);
    
    // 执行推理
    context->enqueueV2(deviceBuffers, stream, nullptr);
    
    // 将输出数据从设备复制回主机
    cudaMemcpyAsync(outputData.data(), deviceBuffers[outputIndex],
                   outputSize * sizeof(float), cudaMemcpyDeviceToHost, stream);
    
    // 同步CUDA流，确保所有操作完成
    cudaStreamSynchronize(stream);
    
    // 输出结果（仅显示前5个值）
    std::cout << "推理结果 (前5个值):" << std::endl;
    for (int i = 0; i < std::min(5, outputSize); i++) {
        std::cout << outputData[i] << " ";
    }
    std::cout << std::endl;
    
    // 清理资源
    cudaFree(deviceBuffers[inputIndex]);
    cudaFree(deviceBuffers[outputIndex]);
    cudaStreamDestroy(stream);
    
    std::cout << "推理完成!" << std::endl;
    return 0;
}

高级C++ API功能

1. 动态形状支持
cpp 复制代码 
// 创建优化配置文件，支持动态输入形状
TRTUniquePtr<nvinfer1::IOptimizationProfile> profile{builder->createOptimizationProfile()};

// 设置输入的最小、最优和最大形状
// 参数：输入名称，选择器类型，维度
profile->setDimensions("input", nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, 3, 224, 224));
profile->setDimensions("input", nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(4, 3, 224, 224));
profile->setDimensions("input", nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(16, 3, 224, 224));

// 将优化配置文件添加到构建配置
config->addOptimizationProfile(profile.get());

// 在推理时设置实际输入形状
context->setBindingDimensions(inputIndex, nvinfer1::Dims4(batchSize, 3, 224, 224));
2. INT8量化
cpp 复制代码 
// 创建INT8校准器（这里使用一个自定义校准器示例）
class MyCalibrator : public nvinfer1::IInt8Calibrator {
    // 实现校准器接口方法
    // ...
};

// 创建校准器实例
MyCalibrator calibrator(/* 校准参数 */);

// 启用INT8模式并设置校准器
if (builder->platformHasFastInt8()) {
    std::cout << "启用INT8精度" << std::endl;
    config->setFlag(nvinfer1::BuilderFlag::kINT8);
    config->setInt8Calibrator(&calibrator);
}
3. 使用CUDA图
cpp 复制代码 
// 创建CUDA图捕获器
cudaGraph_t graph;
cudaGraphExec_t graphExec;

// 开始捕获
cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);

// 执行要捕获的操作
context->enqueueV2(deviceBuffers, stream, nullptr);

// 结束捕获
cudaStreamEndCapture(stream, &graph);

// 实例化图
cudaGraphInstantiate(&graphExec, graph, nullptr, nullptr, 0);

// 使用图执行推理（可以重复执行）
cudaGraphLaunch(graphExec, stream);
cudaStreamSynchronize(stream);

// 清理资源
cudaGraphDestroy(graph);
cudaGraphExecDestroy(graphExec);

Python API详解

Python API提供了与C++ API类似的功能，但使用更简洁的Python语法。

核心类和对象

TensorRT Python API的核心类包括：

  • trt.Logger：日志记录器
  • trt.Builder：构建器
  • trt.NetworkDefinition：网络定义
  • trt.BuilderConfig：构建配置
  • trt.ICudaEngine：CUDA引擎
  • trt.IExecutionContext：执行上下文

详细的Python API示例

下面是一个完整的Python API使用示例，包含详细的中文注释：

python 复制代码 
import os
import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit

class TensorRTModel:
    """TensorRT模型封装类"""
    
    def __init__(self):
        """初始化TensorRT模型"""
        # 创建logger
        self.logger = trt.Logger(trt.Logger.WARNING)
        
        # 创建builder和网络
        self.builder = trt.Builder(self.logger)
        self.network = None
        self.engine = None
        self.context = None
        self.stream = None
        self.host_inputs = None
        self.host_outputs = None
        self.device_inputs = None
        self.device_outputs = None
        self.bindings = None
    
    def build_engine_from_onnx(self, onnx_file_path, engine_file_path=None, precision='fp32'):
        """
        从ONNX模型构建TensorRT引擎
        
        参数:
            onnx_file_path: ONNX模型文件路径
            engine_file_path: 输出的TensorRT引擎文件路径（可选）
            precision: 精度模式，可选 'fp32', 'fp16', 'int8'
        
        返回:
            构建好的TensorRT引擎
        """
        # 创建网络定义，使用显式批处理维度
        EXPLICIT_BATCH = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        self.network = self.builder.create_network(EXPLICIT_BATCH)
        
        # 创建ONNX解析器
        parser = trt.OnnxParser(self.network, self.logger)
        
        # 解析ONNX文件
        print(f"正在解析ONNX模型: {onnx_file_path}")
        with open(onnx_file_path, 'rb') as model:
            if not parser.parse(model.read()):
                print('ERROR: 解析ONNX模型失败')
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                return None
        
        # 创建构建配置
        config = self.builder.create_builder_config()
        
        # 设置最大工作空间大小（1GB）
        config.max_workspace_size = 1 << 30
        
        # 根据指定精度设置标志
        if precision == 'fp16' and self.builder.platform_has_fast_fp16:
            print('启用FP16精度')
            config.set_flag(trt.BuilderFlag.FP16)
        elif precision == 'int8' and self.builder.platform_has_fast_int8:
            print('启用INT8精度')
            config.set_flag(trt.BuilderFlag.INT8)
            # 注意：INT8需要校准器，这里简化处理
        
        # 构建引擎
        print('构建TensorRT引擎...')
        serialized_engine = self.builder.build_serialized_network(self.network, config)
        
        if not serialized_engine:
            print('ERROR: 引擎构建失败')
            return None
        
        # 创建运行时并反序列化引擎
        runtime = trt.Runtime(self.logger)
        self.engine = runtime.deserialize_cuda_engine(serialized_engine)
        
        # 保存引擎到文件（如果指定了文件路径）
        if engine_file_path:
            with open(engine_file_path, 'wb') as f:
                f.write(serialized_engine)
            print(f'引擎已保存到: {engine_file_path}')
        
        return self.engine
    
    def load_engine(self, engine_file_path):
        """
        从文件加载TensorRT引擎
        
        参数:
            engine_file_path: TensorRT引擎文件路径
        
        返回:
            加载的TensorRT引擎
        """
        print(f"从文件加载引擎: {engine_file_path}")
        
        # 读取引擎文件
        with open(engine_file_path, 'rb') as f:
            serialized_engine = f.read()
        
        # 创建运行时并反序列化引擎
        runtime = trt.Runtime(self.logger)
        self.engine = runtime.deserialize_cuda_engine(serialized_engine)
        
        if not self.engine:
            print('ERROR: 引擎加载失败')
            return None
        
        return self.engine
    
    def prepare_inference(self):
        """准备推理环境"""
        if not self.engine:
            print('ERROR: 引擎未初始化')
            return False
        
        # 创建执行上下文
        self.context = self.engine.create_execution_context()
        
        # 创建CUDA流
        self.stream = cuda.Stream()
        
        # 获取绑定索引
        input_binding_idxs = []
        output_binding_idxs = []
        
        for i in range(self.engine.num_bindings):
            if self.engine.binding_is_input(i):
                input_binding_idxs.append(i)
            else:
                output_binding_idxs.append(i)
        
        # 分配输出缓冲区
        self.output_buffers = {}
        
        for i in output_binding_idxs:
            name = self.engine.get_binding_name(i)
            dtype = trt.nptype(self.engine.get_binding_dtype(i))
            shape = self.context.get_binding_shape(i)
            size = trt.volume(shape)
            
            # 分配主机和设备内存
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            
            self.output_buffers[name] = {
                'host': host_mem,
                'device': device_mem,
                'shape': shape
            }
        
        return True
    
    def infer(self, input_data):
        """
        执行推理
        
        参数:
            input_data: 输入数据，可以是numpy数组或字典（键为输入名称，值为numpy数组）
            
        返回:
            输出数据列表
        """
        if not self.engine or not self.context:
            if not self.prepare_inference():
                return None
        
        # 如果输入是numpy数组，假设只有一个输入
        if isinstance(input_data, np.ndarray):
            input_data = {self.engine.get_binding_name(0): input_data}
        
        # 准备输入绑定
        bindings = []
        input_buffers = {}
        
        # 添加输入绑定
        for i in range(self.engine.num_bindings):
            if self.engine.binding_is_input(i):
                name = self.engine.get_binding_name(i)
                
                if name not in input_data:
                    print(f"ERROR: 输入 '{name}' 未提供")
                    return None
                
                data = input_data[name]
                
                # 设置动态形状（如果需要）
                if self.engine.is_shape_binding(i) or -1 in self.context.get_binding_shape(i):
                    self.context.set_binding_shape(i, data.shape)
                
                # 分配主机和设备内存
                dtype = trt.nptype(self.engine.get_binding_dtype(i))
                host_mem = cuda.pagelocked_empty(data.size, dtype)
                device_mem = cuda.mem_alloc(host_mem.nbytes)
                
                # 复制数据
                np.copyto(host_mem, data.ravel())
                cuda.memcpy_htod_async(device_mem, host_mem, self.stream)
                
                input_buffers[name] = {
                    'host': host_mem,
                    'device': device_mem
                }
                
                bindings.append(int(device_mem))
            else:
                name = self.engine.get_binding_name(i)
                bindings.append(int(self.output_buffers[name]['device']))
        
        # 执行推理
        self.context.execute_async_v2(bindings=bindings, stream_handle=self.stream.handle)
        
        # 将输出从设备复制到主机
        outputs = []
        
        for i in range(self.engine.num_bindings):
            if not self.engine.binding_is_input(i):
                name = self.engine.get_binding_name(i)
                shape = self.context.get_binding_shape(i)
                size = trt.volume(shape)
                
                cuda.memcpy_dtoh_async(self.output_buffers[name]['host'], 
                                      self.output_buffers[name]['device'], 
                                      self.stream)
                
                # 同步流
                self.stream.synchronize()
                
                # 重塑输出
                outputs.append(self.output_buffers[name]['host'][:size].reshape(shape))
        
        # 清理输入缓冲区
        for name, buffer in input_buffers.items():
            buffer['device'].free()
        
        return outputs
    
    def __del__(self):
        """清理资源"""
        # 释放设备内存
        if hasattr(self, 'device_inputs') and self.device_inputs:
            for mem in self.device_inputs:
                mem.free()
        
        if hasattr(self, 'device_outputs') and self.device_outputs:
            for mem in self.device_outputs:
                mem.free()

# 使用示例
if __name__ == '__main__':
    # 创建TensorRT模型
    model = TensorRTModel()
    
    # 从ONNX模型构建引擎
    onnx_file = 'resnet50.onnx'
    engine_file = 'resnet50.engine'
    
    # 如果引擎文件已存在，直接加载
    if os.path.exists(engine_file):
        model.load_engine(engine_file)
    else:
        # 否则从ONNX模型构建引擎
        model.build_engine_from_onnx(onnx_file, engine_file, precision='fp16')
    
    # 准备输入数据
    input_shape = (1, 3, 224, 224)  # 批大小=1, 3通道, 224x224图像
    input_data = np.random.random(input_shape).astype(np.float32)
    
    # 执行推理
    outputs = model.infer(input_data)
    
    # 处理输出
    if outputs:
        # 获取前5个最高概率的类别
        probs = outputs[0][0]
        top5_indices = np.argsort(probs)[-5:][::-1]
        
        print("\nTop 5预测结果:")
        for i, idx in enumerate(top5_indices):
            print(f"  {i+1}. 类别 {idx}: {probs[idx]:.6f}")

高级Python API功能

1. 动态形状支持
python 复制代码 
# 创建优化配置文件，支持动态输入形状
profile = builder.create_optimization_profile()

# 设置输入的最小、最优和最大形状
profile.set_shape(
    "input",                          # 输入名称
    min=(1, 3, 224, 224),             # 最小形状
    opt=(4, 3, 224, 224),             # 优化形状
    max=(16, 3, 224, 224)             # 最大形状
)

# 将优化配置文件添加到构建配置
config.add_optimization_profile(profile)

# 在推理时设置实际输入形状
context.set_binding_shape(0, (batch_size, 3, 224, 224))
2. INT8量化与校准
python 复制代码 
# 自定义校准器
class MyCalibrator(trt.IInt8Calibrator):
    def __init__(self, calibration_data, batch_size, cache_file):
        """
        初始化校准器
        
        参数:
            calibration_data: 校准数据集
            batch_size: 批大小
            cache_file: 校准缓存文件路径
        """
        super().__init__()
        self.calibration_data = calibration_data
        self.batch_size = batch_size
        self.cache_file = cache_file
        self.current_index = 0
        self.device_input = None
        
        # 分配设备内存
        input_size = batch_size * 3 * 224 * 224  # 假设输入是224x224的RGB图像
        self.device_input = cuda.mem_alloc(input_size * 4)  # 4字节的float32
    
    def get_batch_size(self):
        """返回批大小"""
        return self.batch_size
    
    def get_batch(self, names):
        """获取下一批校准数据"""
        if self.current_index >= len(self.calibration_data):
            return None
        
        # 获取当前批次数据
        batch = self.calibration_data[self.current_index:self.current_index+self.batch_size]
        self.current_index += self.batch_size
        
        # 如果批次不完整，填充到完整批次
        if len(batch) < self.batch_size:
            batch = np.pad(batch, ((0, self.batch_size - len(batch)), (0, 0), (0, 0), (0, 0)), 'constant')
        
        # 将数据复制到设备
        cuda.memcpy_htod(self.device_input, np.ascontiguousarray(batch))
        
        return [self.device_input]
    
    def read_calibration_cache(self):
        """读取校准缓存"""
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()
        return None
    
    def write_calibration_cache(self, cache):
        """写入校准缓存"""
        with open(self.cache_file, "wb") as f:
            f.write(cache)

# 使用校准器
calibration_data = load_calibration_data()  # 加载校准数据
calibrator = MyCalibrator(calibration_data, batch_size=8, cache_file="calibration.cache")

# 启用INT8并设置校准器
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = calibrator
3. 使用CUDA图
python 复制代码 
# 使用CUDA图加速推理
import pycuda.driver as cuda

# 创建CUDA流
stream = cuda.Stream()

# 开始捕获
cuda.graph_begin_capture(stream, cuda.graph_capture_mode.GLOBAL)

# 执行要捕获的操作
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)

# 结束捕获
graph = cuda.graph_end_capture(stream)

# 实例化图
graph_exec = cuda.graph_instantiate(graph)

# 使用图执行推理（可以重复执行）
graph_exec.launch(stream)
stream.synchronize()

TensorRT API最佳实践

1. 内存管理

  • 使用智能指针：在C++中使用智能指针管理TensorRT对象的生命周期
  • 使用页锁定内存：使用CUDA页锁定内存（pinned memory）减少主机和设备之间的数据传输开销
  • 重用缓冲区：对于多次推理，重用已分配的内存缓冲区

2. 错误处理

  • 检查返回值：TensorRT API的大多数函数都返回状态或指针，应该检查这些返回值
  • 使用日志记录器：自定义ILogger实现，根据需要记录不同级别的消息
  • 解析错误处理：特别是在解析ONNX模型时，检查并处理解析错误

3. 性能优化

  • 批处理：尽可能使用批处理来提高吞吐量
  • 选择合适的精度：根据应用需求选择FP32、FP16或INT8精度
  • 使用CUDA图：对于固定输入大小的推理，使用CUDA图进一步提高性能
  • 优化内存访问模式：减少主机和设备之间的数据传输

4. 可移植性

  • 序列化引擎：将构建好的引擎序列化到文件，避免每次运行时重新构建
  • 检查硬件兼容性：在使用特定功能（如FP16或INT8）前检查硬件支持
  • 处理动态形状：使用优化配置文件支持动态输入形状

总结

本章详细介绍了TensorRT的C++和Python API，包括核心类、基本用法和高级功能。通过丰富的代码示例和中文注释，读者应该能够全面了解如何使用TensorRT API来优化和部署深度学习模型。

在下一章中，我们将探讨TensorRT在各种实际应用场景中的使用方法，包括图像分类、目标检测、语义分割等，并提供完整的端到端示例。

第五部分：TensorRT应用场景实战

在前面的章节中，我们已经详细介绍了TensorRT的基本概念、安装方法、基础使用流程以及API详解。本章将通过几个实际应用场景的完整示例，展示TensorRT在不同深度学习任务中的应用，帮助读者将所学知识应用到实际项目中。

图像分类实战

图像分类是深度学习中最基础也是最常见的应用之一。下面我们将展示如何使用TensorRT优化和部署一个预训练的ResNet-50模型用于图像分类任务。

准备预训练模型

首先，我们需要准备一个预训练的ResNet-50模型并将其转换为ONNX格式：

python 复制代码 
# 准备预训练的ResNet-50模型并转换为ONNX格式
import torch
import torchvision.models as models

def prepare_resnet50_model():
    """
    准备预训练的ResNet-50模型并导出为ONNX格式
    
    返回:
        导出的ONNX模型文件路径
    """
    print("准备预训练的ResNet-50模型...")
    
    # 加载预训练模型
    model = models.resnet50(pretrained=True)
    model.eval()
    
    # 创建示例输入
    dummy_input = torch.randn(1, 3, 224, 224)
    
    # 导出ONNX模型
    onnx_file_path = "resnet50.onnx"
    torch.onnx.export(model,               # 要导出的模型
                     dummy_input,          # 模型输入
                     onnx_file_path,       # 输出文件名
                     export_params=True,   # 存储训练好的参数权重
                     opset_version=13,     # ONNX算子集版本
                     do_constant_folding=True,  # 是否执行常量折叠优化
                     input_names=["input"],     # 输入名称
                     output_names=["output"],   # 输出名称
                     dynamic_axes={"input": {0: "batch_size"},    # 动态轴
                                  "output": {0: "batch_size"}})
    
    print(f"ONNX模型已导出到: {onnx_file_path}")
    return onnx_file_path

构建图像分类应用

接下来，我们创建一个完整的图像分类应用，包括图像预处理、TensorRT推理和结果后处理：

python 复制代码 
import os
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from PIL import Image

class ImageClassifier:
    """使用TensorRT的图像分类器"""
    
    def __init__(self, engine_file_path, labels_file_path):
        """
        初始化图像分类器
        
        参数:
            engine_file_path: TensorRT引擎文件路径
            labels_file_path: 类别标签文件路径
        """
        # 加载类别标签
        self.labels = self._load_labels(labels_file_path)
        
        # 创建TensorRT推理引擎
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        
        # 加载引擎
        with open(engine_file_path, 'rb') as f:
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        
        # 创建执行上下文
        self.context = self.engine.create_execution_context()
        
        # 获取输入和输出的详细信息
        self.input_idx = self.engine.get_binding_index("input")
        self.output_idx = self.engine.get_binding_index("output")
        
        # 获取输入形状
        self.input_shape = self.engine.get_binding_shape(self.input_idx)
        
        # 创建CUDA流
        self.stream = cuda.Stream()
        
        # 分配缓冲区
        self._allocate_buffers()
    
    def _load_labels(self, labels_file_path):
        """
        加载类别标签
        
        参数:
            labels_file_path: 标签文件路径
            
        返回:
            标签列表
        """
        with open(labels_file_path, 'r') as f:
            return [line.strip() for line in f.readlines()]
    
    def _allocate_buffers(self):
        """分配输入和输出缓冲区"""
        # 计算输入和输出大小
        self.input_size = trt.volume(self.input_shape)
        self.output_shape = self.engine.get_binding_shape(self.output_idx)
        self.output_size = trt.volume(self.output_shape)
        
        # 分配主机和设备内存
        self.host_input = cuda.pagelocked_empty(self.input_size, dtype=np.float32)
        self.host_output = cuda.pagelocked_empty(self.output_size, dtype=np.float32)
        
        self.device_input = cuda.mem_alloc(self.host_input.nbytes)
        self.device_output = cuda.mem_alloc(self.host_output.nbytes)
        
        # 准备绑定
        self.bindings = [int(self.device_input), int(self.device_output)]
    
    def preprocess_image(self, image_path):
        """
        预处理图像
        
        参数:
            image_path: 图像文件路径
            
        返回:
            预处理后的图像数据
        """
        # 加载图像
        image = Image.open(image_path).convert('RGB')
        
        # 调整大小
        image = image.resize((224, 224))
        
        # 转换为numpy数组
        image_data = np.array(image).astype(np.float32)
        
        # 调整通道顺序：从HWC到CHW
        image_data = image_data.transpose((2, 0, 1))
        
        # 标准化
        mean = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))
        std = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))
        image_data = (image_data / 255.0 - mean) / std
        
        # 添加批次维度
        image_data = np.expand_dims(image_data, axis=0)
        
        return image_data
    
    def classify(self, image_path):
        """
        对图像进行分类
        
        参数:
            image_path: 图像文件路径
            
        返回:
            分类结果（类别ID、类别名称和置信度）
        """
        # 预处理图像
        input_data = self.preprocess_image(image_path)
        
        # 复制输入数据到主机内存
        np.copyto(self.host_input, input_data.ravel())
        
        # 将输入数据从主机复制到设备
        cuda.memcpy_htod_async(self.device_input, self.host_input, self.stream)
        
        # 执行推理
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        
        # 将输出数据从设备复制回主机
        cuda.memcpy_dtoh_async(self.host_output, self.device_output, self.stream)
        
        # 同步流
        self.stream.synchronize()
        
        # 处理输出
        output = self.host_output.reshape(self.output_shape)
        
        # 获取前5个最高概率的类别
        top5_indices = np.argsort(output[0])[-5:][::-1]
        
        results = []
        for i, idx in enumerate(top5_indices):
            if idx < len(self.labels):
                results.append({
                    'rank': i + 1,
                    'class_id': int(idx),
                    'class_name': self.labels[idx],
                    'confidence': float(output[0][idx])
                })
        
        return results
    
    def __del__(self):
        """清理资源"""
        # 释放设备内存
        if hasattr(self, 'device_input'):
            self.device_input.free()
        
        if hasattr(self, 'device_output'):
            self.device_output.free()

# 使用示例
def main():
    # 准备模型
    onnx_file_path = prepare_resnet50_model()
    
    # 构建TensorRT引擎
    from tensorrt_utils import build_engine_from_onnx
    engine_file_path = "resnet50.engine"
    
    if not os.path.exists(engine_file_path):
        build_engine_from_onnx(onnx_file_path, engine_file_path, precision='fp16')
    
    # 加载ImageNet类别标签
    labels_file_path = "imagenet_classes.txt"
    
    # 创建图像分类器
    classifier = ImageClassifier(engine_file_path, labels_file_path)
    
    # 对图像进行分类
    image_path = "cat.jpg"
    results = classifier.classify(image_path)
    
    # 打印结果
    print(f"\n对图像 {image_path} 的分类结果:")
    for result in results:
        print(f"  {result['rank']}. {result['class_name']} ({result['confidence']:.4f})")

if __name__ == "__main__":
    main()

目标检测实战

目标检测是计算机视觉中的一个重要任务，它不仅需要识别图像中的对象，还需要定位它们的位置。下面我们将展示如何使用TensorRT优化和部署YOLOv5模型进行目标检测。

准备YOLOv5模型

首先，我们需要准备一个预训练的YOLOv5模型并将其转换为ONNX格式：

python 复制代码 
# 准备YOLOv5模型并转换为ONNX格式
import torch

def prepare_yolov5_model():
    """
    准备YOLOv5模型并导出为ONNX格式
    
    返回:
        导出的ONNX模型文件路径
    """
    print("准备YOLOv5模型...")
    
    # 从PyTorch Hub加载预训练模型
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
    model.eval()
    
    # 导出ONNX模型
    onnx_file_path = "yolov5s.onnx"
    
    # 设置导出参数
    img_size = (640, 640)  # 输入图像大小
    batch_size = 1         # 批处理大小
    
    # 导出模型
    torch.onnx.export(model,                      # 要导出的模型
                     torch.zeros(batch_size, 3, *img_size),  # 模型输入
                     onnx_file_path,              # 输出文件名
                     export_params=True,          # 存储训练好的参数权重
                     opset_version=13,            # ONNX算子集版本
                     do_constant_folding=True,    # 是否执行常量折叠优化
                     input_names=["input"],       # 输入名称
                     output_names=["output"],     # 输出名称
                     dynamic_axes={"input": {0: "batch_size"},  # 动态轴
                                  "output": {0: "batch_size"}})
    
    print(f"ONNX模型已导出到: {onnx_file_path}")
    return onnx_file_path

构建目标检测应用

接下来，我们创建一个完整的目标检测应用，包括图像预处理、TensorRT推理和结果后处理：

python 复制代码 
import os
import cv2
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

class ObjectDetector:
    """使用TensorRT的目标检测器"""
    
    def __init__(self, engine_file_path, class_names_file):
        """
        初始化目标检测器
        
        参数:
            engine_file_path: TensorRT引擎文件路径
            class_names_file: 类别名称文件路径
        """
        # 加载类别名称
        self.class_names = self._load_class_names(class_names_file)
        
        # 创建TensorRT推理引擎
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        
        # 加载引擎
        with open(engine_file_path, 'rb') as f:
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        
        # 创建执行上下文
        self.context = self.engine.create_execution_context()
        
        # 获取输入和输出的详细信息
        self.input_idx = self.engine.get_binding_index("input")
        self.output_idx = self.engine.get_binding_index("output")
        
        # 获取输入形状
        self.input_shape = self.engine.get_binding_shape(self.input_idx)
        self.img_size = (self.input_shape[2], self.input_shape[3])  # HW
        
        # 创建CUDA流
        self.stream = cuda.Stream()
        
        # 分配缓冲区
        self._allocate_buffers()
        
        # 设置NMS参数
        self.conf_threshold = 0.25
        self.iou_threshold = 0.45
    
    def _load_class_names(self, class_names_file):
        """
        加载类别名称
        
        参数:
            class_names_file: 类别名称文件路径
            
        返回:
            类别名称列表
        """
        with open(class_names_file, 'r') as f:
            return [line.strip() for line in f.readlines()]
    
    def _allocate_buffers(self):
        """分配输入和输出缓冲区"""
        # 计算输入和输出大小
        self.input_size = trt.volume(self.input_shape)
        self.output_shape = self.engine.get_binding_shape(self.output_idx)
        self.output_size = trt.volume(self.output_shape)
        
        # 分配主机和设备内存
        self.host_input = cuda.pagelocked_empty(self.input_size, dtype=np.float32)
        self.host_output = cuda.pagelocked_empty(self.output_size, dtype=np.float32)
        
        self.device_input = cuda.mem_alloc(self.host_input.nbytes)
        self.device_output = cuda.mem_alloc(self.host_output.nbytes)
        
        # 准备绑定
        self.bindings = [int(self.device_input), int(self.device_output)]
    
    def preprocess_image(self, image_path):
        """
        预处理图像
        
        参数:
            image_path: 图像文件路径
            
        返回:
            预处理后的图像数据和原始图像信息
        """
        # 读取图像
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图像: {image_path}")
        
        # 保存原始图像信息
        original_img = img.copy()
        original_shape = img.shape
        
        # 调整大小并保持宽高比
        h, w = original_shape[:2]
        scale = min(self.img_size[0] / h, self.img_size[1] / w)
        new_h, new_w = int(h * scale), int(w * scale)
        
        # 调整图像大小
        img = cv2.resize(img, (new_w, new_h))
        
        # 创建填充图像
        padded_img = np.zeros((self.img_size[0], self.img_size[1], 3), dtype=np.uint8)
        padded_img[:new_h, :new_w] = img
        
        # 转换为RGB并归一化
        padded_img = cv2.cvtColor(padded_img, cv2.COLOR_BGR2RGB)
        padded_img = padded_img.astype(np.float32) / 255.0
        
        # 调整通道顺序：从HWC到CHW
        padded_img = padded_img.transpose((2, 0, 1))
        
        # 添加批次维度
        padded_img = np.expand_dims(padded_img, axis=0)
        
        return padded_img, original_img, scale, (new_h, new_w)
    
    def postprocess(self, output, original_img, scale, new_shape):
        """
        后处理检测结果
        
        参数:
            output: 模型输出
            original_img: 原始图像
            scale: 缩放比例
            new_shape: 调整大小后的形状
            
        返回:
            处理后的检测结果
        """
        # 重塑输出
        output = output.reshape((-1, self.output_shape[-1]))
        
        # 应用置信度阈值
        valid = output[:, 4] > self.conf_threshold
        output = output[valid]
        
        if len(output) == 0:
            return []
        
        # 分离坐标、置信度和类别概率
        boxes = output[:, :4]
        scores = output[:, 4]
        classes = output[:, 5:]
        
        # 获取类别ID和置信度
        class_ids = np.argmax(classes, axis=1)
        confidences = scores * np.max(classes, axis=1)
        
        # 应用置信度阈值
        valid = confidences > self.conf_threshold
        boxes = boxes[valid]
        class_ids = class_ids[valid]
        confidences = confidences[valid]
        
        if len(boxes) == 0:
            return []
        
        # 转换坐标格式：从[x_center, y_center, width, height]到[x1, y1, x2, y2]
        x_center = boxes[:, 0]
        y_center = boxes[:, 1]
        width = boxes[:, 2]
        height = boxes[:, 3]
        
        x1 = x_center - width / 2
        y1 = y_center - height / 2
        x2 = x_center + width / 2
        y2 = y_center + height / 2
        
        # 调整坐标到原始图像
        new_h, new_w = new_shape
        x1 = x1 * new_w / scale
        y1 = y1 * new_h / scale
        x2 = x2 * new_w / scale
        y2 = y2 * new_h / scale
        
        # 应用非极大值抑制
        indices = cv2.dnn.NMSBoxes(
            np.column_stack((x1, y1, x2 - x1, y2 - y1)).tolist(),
            confidences.tolist(),
            self.conf_threshold,
            self.iou_threshold
        )
        
        # 准备结果
        results = []
        for idx in indices:
            if isinstance(idx, list):  # OpenCV 4.5.4之前的版本
                idx = idx[0]
            
            x1_val, y1_val = max(0, int(x1[idx])), max(0, int(y1[idx]))
            x2_val, y2_val = min(original_img.shape[1], int(x2[idx])), min(original_img.shape[0], int(y2[idx]))
            
            class_id = int(class_ids[idx])
            class_name = self.class_names[class_id] if class_id < len(self.class_names) else f"Unknown-{class_id}"
            confidence = float(confidences[idx])
            
            results.append({
                'class_id': class_id,
                'class_name': class_name,
                'confidence': confidence,
                'box': [x1_val, y1_val, x2_val, y2_val]
            })
        
        return results
    
    def detect(self, image_path):
        """
        检测图像中的对象
        
        参数:
            image_path: 图像文件路径
            
        返回:
            检测结果和可视化图像
        """
        # 预处理图像
        input_data, original_img, scale, new_shape = self.preprocess_image(image_path)
        
        # 复制输入数据到主机内存
        np.copyto(self.host_input, input_data.ravel())
        
        # 将输入数据从主机复制到设备
        cuda.memcpy_htod_async(self.device_input, self.host_input, self.stream)
        
        # 执行推理
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        
        # 将输出数据从设备复制回主机
        cuda.memcpy_dtoh_async(self.host_output, self.device_output, self.stream)
        
        # 同步流
        self.stream.synchronize()
        
        # 后处理检测结果
        results = self.postprocess(self.host_output, original_img, scale, new_shape)
        
        # 可视化结果
        vis_img = original_img.copy()
        for result in results:
            # 获取边界框和类别信息
            x1, y1, x2, y2 = result['box']
            class_name = result['class_name']
            confidence = result['confidence']
            
            # 绘制边界框
            color = (0, 255, 0)  # 绿色
            cv2.rectangle(vis_img, (x1, y1), (x2, y2), color, 2)
            
            # 绘制标签
            label = f"{class_name} {confidence:.2f}"
            (label_width, label_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            cv2.rectangle(vis_img, (x1, y1 - label_height - baseline), (x1 + label_width, y1), color, -1)
            cv2.putText(vis_img, label, (x1, y1 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
        
        return results, vis_img
    
    def __del__(self):
        """清理资源"""
        # 释放设备内存
        if hasattr(self, 'device_input'):
            self.device_input.free()
        
        if hasattr(self, 'device_output'):
            self.device_output.free()

# 使用示例
def main():
    # 准备模型
    onnx_file_path = prepare_yolov5_model()
    
    # 构建TensorRT引擎
    from tensorrt_utils import build_engine_from_onnx
    engine_file_path = "yolov5s.engine"
    
    if not os.path.exists(engine_file_path):
        build_engine_from_onnx(onnx_file_path, engine_file_path, precision='fp16')
    
    # 加载COCO类别名称
    class_names_file = "coco_classes.txt"
    
    # 创建目标检测器
    detector = ObjectDetector(engine_file_path, class_names_file)
    
    # 对图像进行检测
    image_path = "street.jpg"
    results, vis_img = detector.detect(image_path)
    
    # 保存可视化结果
    output_path = "detection_result.jpg"
    cv2.imwrite(output_path, vis_img)
    
    # 打印结果
    print(f"\n对图像 {image_path} 的检测结果:")
    for i, result in enumerate(results):
        print(f"  {i+1}. {result['class_name']} ({result['confidence']:.4f}) at {result['box']}")
    print(f"\n可视化结果已保存到: {output_path}")

if __name__ == "__main__":
    main()

语义分割实战

语义分割是像素级的分类任务，为图像中的每个像素分配一个类别标签。下面我们将展示如何使用TensorRT优化和部署DeepLabV3模型进行语义分割。

准备DeepLabV3模型

首先，我们需要准备一个预训练的DeepLabV3模型并将其转换为ONNX格式：

python 复制代码 
# 准备DeepLabV3模型并转换为ONNX格式
import torch
import torchvision.models.segmentation as models

def prepare_deeplabv3_model():
    """
    准备DeepLabV3模型并导出为ONNX格式
    
    返回:
        导出的ONNX模型文件路径
    """
    print("准备DeepLabV3模型...")
    
    # 加载预训练模型
    model = models.deeplabv3_resnet50(pretrained=True)
    model.eval()
    
    # 创建示例输入
    dummy_input = torch.randn(1, 3, 513, 513)
    
    # 导出ONNX模型
    onnx_file_path = "deeplabv3.onnx"
    torch.onnx.export(model,               # 要导出的模型
                     dummy_input,          # 模型输入
                     onnx_file_path,       # 输出文件名
                     export_params=True,   # 存储训练好的参数权重
                     opset_version=13,     # ONNX算子集版本
                     do_constant_folding=True,  # 是否执行常量折叠优化
                     input_names=["input"],     # 输入名称
                     output_names=["output"],   # 输出名称
                     dynamic_axes={"input": {0: "batch_size"},    # 动态轴
                                  "output": {0: "batch_size"}})
    
    print(f"ONNX模型已导出到: {onnx_file_path}")
    return onnx_file_path

构建语义分割应用

接下来，我们创建一个完整的语义分割应用，包括图像预处理、TensorRT推理和结果可视化：

python 复制代码 
import os
import cv2
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from PIL import Image

class SemanticSegmenter:
    """使用TensorRT的语义分割器"""
    
    def __init__(self, engine_file_path, num_classes=21):
        """
        初始化语义分割器
        
        参数:
            engine_file_path: TensorRT引擎文件路径
            num_classes: 类别数量（默认为21，对应Pascal VOC数据集）
        """
        self.num_classes = num_classes
        
        # 创建TensorRT推理引擎
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        
        # 加载引擎
        with open(engine_file_path, 'rb') as f:
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        
        # 创建执行上下文
        self.context = self.engine.create_execution_context()
        
        # 获取输入和输出的详细信息
        self.input_idx = self.engine.get_binding_index("input")
        self.output_idx = self.engine.get_binding_index("output")
        
        # 获取输入形状
        self.input_shape = self.engine.get_binding_shape(self.input_idx)
        self.img_size = (self.input_shape[2], self.input_shape[3])  # HW
        
        # 创建CUDA流
        self.stream = cuda.Stream()
        
        # 分配缓冲区
        self._allocate_buffers()
        
        # 创建颜色映射
        self.color_map = self._create_pascal_label_colormap()
    
    def _create_pascal_label_colormap(self):
        """
        创建Pascal VOC数据集的颜色映射
        
        返回:
            颜色映射数组
        """
        colormap = np.zeros((self.num_classes, 3), dtype=np.uint8)
        for i in range(self.num_classes):
            r = i * 10 % 255
            g = (i * 100) % 255
            b = (i * 29) % 255
            colormap[i] = np.array([r, g, b])
        return colormap
    
    def _allocate_buffers(self):
        """分配输入和输出缓冲区"""
        # 计算输入和输出大小
        self.input_size = trt.volume(self.input_shape)
        self.output_shape = self.engine.get_binding_shape(self.output_idx)
        self.output_size = trt.volume(self.output_shape)
        
        # 分配主机和设备内存
        self.host_input = cuda.pagelocked_empty(self.input_size, dtype=np.float32)
        self.host_output = cuda.pagelocked_empty(self.output_size, dtype=np.float32)
        
        self.device_input = cuda.mem_alloc(self.host_input.nbytes)
        self.device_output = cuda.mem_alloc(self.host_output.nbytes)
        
        # 准备绑定
        self.bindings = [int(self.device_input), int(self.device_output)]
    
    def preprocess_image(self, image_path):
        """
        预处理图像
        
        参数:
            image_path: 图像文件路径
            
        返回:
            预处理后的图像数据和原始图像
        """
        # 读取图像
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图像: {image_path}")
        
        # 保存原始图像
        original_img = img.copy()
        
        # 调整大小
        img = cv2.resize(img, (self.img_size[1], self.img_size[0]))
        
        # 转换为RGB并归一化
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) / 255.0
        
        # 标准化
        mean = np.array([0.485, 0.456, 0.406]).reshape((1, 1, 3))
        std = np.array([0.229, 0.224, 0.225]).reshape((1, 1, 3))
        img = (img - mean) / std
        
        # 调整通道顺序：从HWC到CHW
        img = img.transpose((2, 0, 1))
        
        # 添加批次维度
        img = np.expand_dims(img, axis=0)
        
        return img, original_img
    
    def segment(self, image_path):
        """
        对图像进行语义分割
        
        参数:
            image_path: 图像文件路径
            
        返回:
            分割结果和可视化图像
        """
        # 预处理图像
        input_data, original_img = self.preprocess_image(image_path)
        
        # 复制输入数据到主机内存
        np.copyto(self.host_input, input_data.ravel())
        
        # 将输入数据从主机复制到设备
        cuda.memcpy_htod_async(self.device_input, self.host_input, self.stream)
        
        # 执行推理
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        
        # 将输出数据从设备复制回主机
        cuda.memcpy_dtoh_async(self.host_output, self.device_output, self.stream)
        
        # 同步流
        self.stream.synchronize()
        
        # 处理输出
        output = self.host_output.reshape(self.output_shape)
        
        # 获取分割掩码
        output = output[0]['out'][0]  # 假设输出是字典，包含'out'键
        mask = np.argmax(output, axis=0)
        
        # 调整掩码大小以匹配原始图像
        mask = cv2.resize(mask.astype(np.uint8), (original_img.shape[1], original_img.shape[0]), 
                         interpolation=cv2.INTER_NEAREST)
        
        # 创建可视化图像
        vis_img = self._create_colored_mask(mask, original_img)
        
        return mask, vis_img
    
    def _create_colored_mask(self, mask, original_img, alpha=0.5):
        """
        创建彩色分割掩码
        
        参数:
            mask: 分割掩码
            original_img: 原始图像
            alpha: 混合透明度
            
        返回:
            彩色分割可视化
        """
        # 创建彩色掩码
        colored_mask = np.zeros_like(original_img)
        for i in range(self.num_classes):
            colored_mask[mask == i] = self.color_map[i]
        
        # 混合原始图像和彩色掩码
        vis_img = cv2.addWeighted(original_img, 1 - alpha, colored_mask, alpha, 0)
        
        return vis_img
    
    def __del__(self):
        """清理资源"""
        # 释放设备内存
        if hasattr(self, 'device_input'):
            self.device_input.free()
        
        if hasattr(self, 'device_output'):
            self.device_output.free()

# 使用示例
def main():
    # 准备模型
    onnx_file_path = prepare_deeplabv3_model()
    
    # 构建TensorRT引擎
    from tensorrt_utils import build_engine_from_onnx
    engine_file_path = "deeplabv3.engine"
    
    if not os.path.exists(engine_file_path):
        build_engine_from_onnx(onnx_file_path, engine_file_path, precision='fp16')
    
    # 创建语义分割器
    segmenter = SemanticSegmenter(engine_file_path)
    
    # 对图像进行分割
    image_path = "street.jpg"
    mask, vis_img = segmenter.segment(image_path)
    
    # 保存可视化结果
    output_path = "segmentation_result.jpg"
    cv2.imwrite(output_path, vis_img)
    
    print(f"\n对图像 {image_path} 的分割结果已保存到: {output_path}")

if __name__ == "__main__":
    main()

文本分类实战

自然语言处理（NLP）是深度学习的另一个重要应用领域。下面我们将展示如何使用TensorRT优化和部署BERT模型进行文本分类。

准备BERT模型

首先，我们需要准备一个预训练的BERT模型并将其转换为ONNX格式：

python 复制代码 
# 准备BERT模型并转换为ONNX格式
import torch
from transformers import BertTokenizer, BertForSequenceClassification

def prepare_bert_model():
    """
    准备BERT模型并导出为ONNX格式
    
    返回:
        导出的ONNX模型文件路径和分词器
    """
    print("准备BERT模型...")
    
    # 加载预训练模型和分词器
    model_name = "bert-base-uncased"
    tokenizer = BertTokenizer.from_pretrained(model_name)
    model = BertForSequenceClassification.from_pretrained(model_name)
    model.eval()
    
    # 创建示例输入
    text = "TensorRT is a high-performance deep learning inference optimizer and runtime."
    encoded_input = tokenizer(text, padding='max_length', max_length=128, truncation=True, return_tensors='pt')
    
    # 导出ONNX模型
    onnx_file_path = "bert.onnx"
    torch.onnx.export(model,                       # 要导出的模型
                     (encoded_input['input_ids'], encoded_input['attention_mask']),  # 模型输入
                     onnx_file_path,               # 输出文件名
                     export_params=True,           # 存储训练好的参数权重
                     opset_version=13,             # ONNX算子集版本
                     do_constant_folding=True,     # 是否执行常量折叠优化
                     input_names=["input_ids", "attention_mask"],  # 输入名称
                     output_names=["output"],      # 输出名称
                     dynamic_axes={"input_ids": {0: "batch_size"},  # 动态轴
                                  "attention_mask": {0: "batch_size"},
                                  "output": {0: "batch_size"}})
    
    print(f"ONNX模型已导出到: {onnx_file_path}")
    return onnx_file_path, tokenizer

构建文本分类应用

接下来，我们创建一个完整的文本分类应用，包括文本预处理、TensorRT推理和结果后处理：

python 复制代码 
import os
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

class BERTClassifier:
    """使用TensorRT的BERT文本分类器"""
    
    def __init__(self, engine_file_path, tokenizer):
        """
        初始化BERT文本分类器
        
        参数:
            engine_file_path: TensorRT引擎文件路径
            tokenizer: BERT分词器
        """
        self.tokenizer = tokenizer
        
        # 创建TensorRT推理引擎
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        
        # 加载引擎
        with open(engine_file_path, 'rb') as f:
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        
        # 创建执行上下文
        self.context = self.engine.create_execution_context()
        
        # 获取输入和输出的详细信息
        self.input_ids_idx = self.engine.get_binding_index("input_ids")
        self.attention_mask_idx = self.engine.get_binding_index("attention_mask")
        self.output_idx = self.engine.get_binding_index("output")
        
        # 获取输入形状
        self.input_ids_shape = self.engine.get_binding_shape(self.input_ids_idx)
        self.attention_mask_shape = self.engine.get_binding_shape(self.attention_mask_idx)
        self.output_shape = self.engine.get_binding_shape(self.output_idx)
        
        # 创建CUDA流
        self.stream = cuda.Stream()
        
        # 分配缓冲区
        self._allocate_buffers()
    
    def _allocate_buffers(self):
        """分配输入和输出缓冲区"""
        # 计算输入和输出大小
        self.input_ids_size = trt.volume(self.input_ids_shape)
        self.attention_mask_size = trt.volume(self.attention_mask_shape)
        self.output_size = trt.volume(self.output_shape)
        
        # 分配主机和设备内存
        self.host_input_ids = cuda.pagelocked_empty(self.input_ids_size, dtype=np.int32)
        self.host_attention_mask = cuda.pagelocked_empty(self.attention_mask_size, dtype=np.int32)
        self.host_output = cuda.pagelocked_empty(self.output_size, dtype=np.float32)
        
        self.device_input_ids = cuda.mem_alloc(self.host_input_ids.nbytes)
        self.device_attention_mask = cuda.mem_alloc(self.host_attention_mask.nbytes)
        self.device_output = cuda.mem_alloc(self.host_output.nbytes)
        
        # 准备绑定
        self.bindings = [int(self.device_input_ids), int(self.device_attention_mask), int(self.device_output)]
    
    def preprocess_text(self, text, max_length=128):
        """
        预处理文本
        
        参数:
            text: 输入文本
            max_length: 最大序列长度
            
        返回:
            预处理后的输入ID和注意力掩码
        """
        # 使用分词器处理文本
        encoded_input = self.tokenizer(text, padding='max_length', max_length=max_length, 
                                      truncation=True, return_tensors='np')
        
        return encoded_input['input_ids'], encoded_input['attention_mask']
    
    def classify(self, text):
        """
        对文本进行分类
        
        参数:
            text: 输入文本
            
        返回:
            分类结果（类别ID和置信度）
        """
        # 预处理文本
        input_ids, attention_mask = self.preprocess_text(text)
        
        # 复制输入数据到主机内存
        np.copyto(self.host_input_ids, input_ids.ravel())
        np.copyto(self.host_attention_mask, attention_mask.ravel())
        
        # 将输入数据从主机复制到设备
        cuda.memcpy_htod_async(self.device_input_ids, self.host_input_ids, self.stream)
        cuda.memcpy_htod_async(self.device_attention_mask, self.host_attention_mask, self.stream)
        
        # 执行推理
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        
        # 将输出数据从设备复制回主机
        cuda.memcpy_dtoh_async(self.host_output, self.device_output, self.stream)
        
        # 同步流
        self.stream.synchronize()
        
        # 处理输出
        output = self.host_output.reshape(self.output_shape)
        
        # 应用softmax
        exp_output = np.exp(output - np.max(output, axis=1, keepdims=True))
        softmax_output = exp_output / np.sum(exp_output, axis=1, keepdims=True)
        
        # 获取预测类别和置信度
        class_id = np.argmax(softmax_output, axis=1)[0]
        confidence = softmax_output[0, class_id]
        
        return class_id, confidence
    
    def __del__(self):
        """清理资源"""
        # 释放设备内存
        if hasattr(self, 'device_input_ids'):
            self.device_input_ids.free()
        
        if hasattr(self, 'device_attention_mask'):
            self.device_attention_mask.free()
        
        if hasattr(self, 'device_output'):
            self.device_output.free()

# 使用示例
def main():
    # 准备模型
    onnx_file_path, tokenizer = prepare_bert_model()
    
    # 构建TensorRT引擎
    from tensorrt_utils import build_engine_from_onnx
    engine_file_path = "bert.engine"
    
    if not os.path.exists(engine_file_path):
        build_engine_from_onnx(onnx_file_path, engine_file_path, precision='fp16')
    
    # 创建文本分类器
    classifier = BERTClassifier(engine_file_path, tokenizer)
    
    # 对文本进行分类
    text = "TensorRT is a high-performance deep learning inference optimizer and runtime."
    class_id, confidence = classifier.classify(text)
    
    # 打印结果
    print(f"\n对文本的分类结果:")
    print(f"  类别ID: {class_id}")
    print(f"  置信度: {confidence:.4f}")

if __name__ == "__main__":
    main()

总结

本章通过四个实际应用场景的完整示例，展示了如何使用TensorRT优化和部署不同类型的深度学习模型：

  1. 图像分类：使用ResNet-50模型对图像进行分类，展示了基本的图像预处理、TensorRT推理和结果后处理流程。
  2. 目标检测：使用YOLOv5模型进行目标检测，展示了如何处理更复杂的模型输出和可视化检测结果。
  3. 语义分割：使用DeepLabV3模型进行语义分割，展示了如何处理像素级预测任务和创建分割可视化。
  4. 文本分类：使用BERT模型进行文本分类，展示了TensorRT在自然语言处理任务中的应用。

这些示例涵盖了TensorRT在不同应用领域的使用方法，包括模型准备、引擎构建、推理执行和结果处理等关键步骤。通过这些实战案例，读者应该能够掌握如何将TensorRT应用到自己的实际项目中，充分发挥其性能优势。

在下一章中，我们将探讨TensorRT的高级特性和优化技巧，帮助读者进一步提升模型的推理性能。

第六部分：TensorRT高级特性与优化技巧

在前面的章节中，我们已经介绍了TensorRT的基本概念、安装方法、基础使用流程、API详解以及实际应用场景。本章将深入探讨TensorRT的高级特性和优化技巧，帮助读者进一步提升模型的推理性能。

量化技术

量化是TensorRT中最重要的优化技术之一，它通过降低计算精度来提高推理速度和减少内存占用。TensorRT支持多种量化精度，包括FP32（单精度浮点）、FP16（半精度浮点）和INT8（8位整数）。

FP16量化

FP16量化是最简单的量化方法，它将模型的权重和激活值从FP32转换为FP16，可以在不显著影响精度的情况下将内存使用量减少一半，并在支持FP16的GPU上获得更高的吞吐量。

python 复制代码 
# 使用FP16精度构建TensorRT引擎
import tensorrt as trt

def build_fp16_engine(onnx_file_path, engine_file_path):
    """
    使用FP16精度构建TensorRT引擎
    
    参数:
        onnx_file_path: ONNX模型文件路径
        engine_file_path: 输出的TensorRT引擎文件路径
    """
    # 创建logger和builder
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    
    # 创建网络定义
    network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    network = builder.create_network(network_flags)
    
    # 创建ONNX解析器
    parser = trt.OnnxParser(network, logger)
    
    # 解析ONNX模型
    with open(onnx_file_path, 'rb') as model:
        if not parser.parse(model.read()):
            print("ERROR: Failed to parse the ONNX file.")
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return False
    
    # 配置builder
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 30  # 1GB
    
    # 启用FP16精度
    if builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
        print("FP16模式已启用")
    else:
        print("平台不支持FP16")
    
    # 构建引擎
    print("构建TensorRT引擎...")
    engine = builder.build_engine(network, config)
    
    # 序列化引擎并保存到文件
    with open(engine_file_path, 'wb') as f:
        f.write(engine.serialize())
    
    print(f"TensorRT引擎已保存到: {engine_file_path}")
    return True

INT8量化

INT8量化可以进一步减少内存使用量和提高推理速度，但需要校准数据来确定最佳的量化参数。TensorRT提供了多种校准方法，包括熵校准、熵校准2和遗留校准。

python 复制代码 
# 使用INT8精度和校准数据构建TensorRT引擎
import tensorrt as trt
import numpy as np
import os
from PIL import Image
import calibrator

class ImageBatchStream:
    """
    图像批次流，用于INT8校准
    """
    def __init__(self, calibration_files, batch_size, input_shape, preprocess_func):
        """
        初始化图像批次流
        
        参数:
            calibration_files: 校准图像文件列表
            batch_size: 批次大小
            input_shape: 输入形状 (C, H, W)
            preprocess_func: 图像预处理函数
        """
        self.batch_size = batch_size
        self.input_shape = input_shape
        self.calibration_files = calibration_files
        self.preprocess_func = preprocess_func
        self.batch_idx = 0
        self.max_batches = len(calibration_files) // batch_size
    
    def reset(self):
        """重置批次索引"""
        self.batch_idx = 0
    
    def next_batch(self):
        """获取下一批次数据"""
        if self.batch_idx < self.max_batches:
            batch_files = self.calibration_files[self.batch_idx * self.batch_size:
                                               (self.batch_idx + 1) * self.batch_size]
            
            batch_data = np.zeros((self.batch_size, *self.input_shape), dtype=np.float32)
            
            for i, file_path in enumerate(batch_files):
                img = Image.open(file_path).convert('RGB')
                batch_data[i] = self.preprocess_func(img)
            
            self.batch_idx += 1
            return np.ascontiguousarray(batch_data)
        else:
            return None

class ImageEntropyCalibrator(trt.IInt8EntropyCalibrator2):
    """
    INT8熵校准器
    """
    def __init__(self, batch_stream, cache_file):
        """
        初始化校准器
        
        参数:
            batch_stream: 批次数据流
            cache_file: 校准缓存文件路径
        """
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.batch_stream = batch_stream
        self.cache_file = cache_file
        self.device_input = None
        
        # 分配设备内存
        input_size = trt.volume(batch_stream.input_shape) * batch_stream.batch_size
        self.device_input = cuda.mem_alloc(input_size * np.dtype(np.float32).itemsize)
    
    def get_batch_size(self):
        """返回批次大小"""
        return self.batch_stream.batch_size
    
    def get_batch(self, names):
        """获取下一批次数据"""
        batch = self.batch_stream.next_batch()
        if batch is None:
            return None
        
        # 将数据复制到设备
        cuda.memcpy_htod(self.device_input, batch)
        return [int(self.device_input)]
    
    def read_calibration_cache(self):
        """读取校准缓存"""
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()
        return None
    
    def write_calibration_cache(self, cache):
        """写入校准缓存"""
        with open(self.cache_file, "wb") as f:
            f.write(cache)

def build_int8_engine(onnx_file_path, engine_file_path, calibration_files, input_shape, preprocess_func):
    """
    使用INT8精度和校准数据构建TensorRT引擎
    
    参数:
        onnx_file_path: ONNX模型文件路径
        engine_file_path: 输出的TensorRT引擎文件路径
        calibration_files: 校准图像文件列表
        input_shape: 输入形状 (C, H, W)
        preprocess_func: 图像预处理函数
    """
    # 创建logger和builder
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    
    # 创建网络定义
    network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    network = builder.create_network(network_flags)
    
    # 创建ONNX解析器
    parser = trt.OnnxParser(network, logger)
    
    # 解析ONNX模型
    with open(onnx_file_path, 'rb') as model:
        if not parser.parse(model.read()):
            print("ERROR: Failed to parse the ONNX file.")
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return False
    
    # 配置builder
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 30  # 1GB
    
    # 启用INT8精度
    if builder.platform_has_fast_int8:
        config.set_flag(trt.BuilderFlag.INT8)
        print("INT8模式已启用")
        
        # 创建校准器
        batch_size = 1
        cache_file = "calibration.cache"
        batch_stream = ImageBatchStream(calibration_files, batch_size, input_shape, preprocess_func)
        calibrator = ImageEntropyCalibrator(batch_stream, cache_file)
        
        # 设置校准器
        config.int8_calibrator = calibrator
    else:
        print("平台不支持INT8")
        return False
    
    # 构建引擎
    print("构建TensorRT引擎...")
    engine = builder.build_engine(network, config)
    
    # 序列化引擎并保存到文件
    with open(engine_file_path, 'wb') as f:
        f.write(engine.serialize())
    
    print(f"TensorRT引擎已保存到: {engine_file_path}")
    return True

动态形状支持

TensorRT支持动态输入形状，这对于处理可变大小的输入（如不同分辨率的图像或不同长度的序列）非常有用。

python 复制代码 
# 使用动态形状构建TensorRT引擎
import tensorrt as trt

def build_dynamic_shape_engine(onnx_file_path, engine_file_path, min_shape, opt_shape, max_shape):
    """
    使用动态形状构建TensorRT引擎
    
    参数:
        onnx_file_path: ONNX模型文件路径
        engine_file_path: 输出的TensorRT引擎文件路径
        min_shape: 最小输入形状 (batch_size, channels, height, width)
        opt_shape: 最优输入形状
        max_shape: 最大输入形状
    """
    # 创建logger和builder
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    
    # 创建网络定义
    network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    network = builder.create_network(network_flags)
    
    # 创建ONNX解析器
    parser = trt.OnnxParser(network, logger)
    
    # 解析ONNX模型
    with open(onnx_file_path, 'rb') as model:
        if not parser.parse(model.read()):
            print("ERROR: Failed to parse the ONNX file.")
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return False
    
    # 配置builder
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 30  # 1GB
    
    # 创建优化配置文件
    profile = builder.create_optimization_profile()
    
    # 设置动态形状
    input_name = "input"  # 输入张量的名称
    profile.set_shape(input_name, min_shape, opt_shape, max_shape)
    config.add_optimization_profile(profile)
    
    # 构建引擎
    print("构建TensorRT引擎...")
    engine = builder.build_engine(network, config)
    
    # 序列化引擎并保存到文件
    with open(engine_file_path, 'wb') as f:
        f.write(engine.serialize())
    
    print(f"TensorRT引擎已保存到: {engine_file_path}")
    return True

使用动态形状的引擎进行推理：

python 复制代码 
# 使用动态形状的引擎进行推理
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

class DynamicShapeInferencer:
    """使用动态形状的TensorRT推理器"""
    
    def __init__(self, engine_file_path):
        """
        初始化推理器
        
        参数:
            engine_file_path: TensorRT引擎文件路径
        """
        # 创建TensorRT推理引擎
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        
        # 加载引擎
        with open(engine_file_path, 'rb') as f:
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        
        # 创建执行上下文
        self.context = self.engine.create_execution_context()
        
        # 获取输入和输出的详细信息
        self.input_idx = self.engine.get_binding_index("input")
        self.output_idx = self.engine.get_binding_index("output")
        
        # 创建CUDA流
        self.stream = cuda.Stream()
    
    def infer(self, input_data):
        """
        执行推理
        
        参数:
            input_data: 输入数据，形状为(batch_size, channels, height, width)
            
        返回:
            推理结果
        """
        # 获取输入形状
        input_shape = input_data.shape
        
        # 设置执行上下文的输入形状
        self.context.set_binding_shape(self.input_idx, input_shape)
        
        # 获取输出形状
        output_shape = self.context.get_binding_shape(self.output_idx)
        
        # 分配主机和设备内存
        host_input = cuda.pagelocked_empty(trt.volume(input_shape), dtype=np.float32)
        host_output = cuda.pagelocked_empty(trt.volume(output_shape), dtype=np.float32)
        
        device_input = cuda.mem_alloc(host_input.nbytes)
        device_output = cuda.mem_alloc(host_output.nbytes)
        
        # 准备绑定
        bindings = [int(device_input), int(device_output)]
        
        # 复制输入数据到主机内存
        np.copyto(host_input, input_data.ravel())
        
        # 将输入数据从主机复制到设备
        cuda.memcpy_htod_async(device_input, host_input, self.stream)
        
        # 执行推理
        self.context.execute_async_v2(bindings=bindings, stream_handle=self.stream.handle)
        
        # 将输出数据从设备复制回主机
        cuda.memcpy_dtoh_async(host_output, device_output, self.stream)
        
        # 同步流
        self.stream.synchronize()
        
        # 释放设备内存
        device_input.free()
        device_output.free()
        
        # 返回结果
        return host_output.reshape(output_shape)

层间融合和张量核心加速

TensorRT会自动执行多种优化，包括层间融合和使用Tensor Cores（如果可用）来加速计算。

层间融合

层间融合是TensorRT的一项重要优化技术，它将多个相邻的层合并为一个层，减少内存访问和计算开销。例如，卷积层、批量归一化层和ReLU激活层可以融合为一个单一的层。

以下是一个示例，展示如何检查TensorRT执行的层间融合：

python 复制代码 
# 检查TensorRT执行的层间融合
import tensorrt as trt

def inspect_engine_layers(engine_file_path):
    """
    检查TensorRT引擎中的层
    
    参数:
        engine_file_path: TensorRT引擎文件路径
    """
    # 创建logger和runtime
    logger = trt.Logger(trt.Logger.WARNING)
    runtime = trt.Runtime(logger)
    
    # 加载引擎
    with open(engine_file_path, 'rb') as f:
        engine = runtime.deserialize_cuda_engine(f.read())
    
    # 获取层信息
    num_layers = engine.num_layers
    print(f"引擎中的层数: {num_layers}")
    
    # 创建检查器
    inspector = engine.create_engine_inspector()
    
    # 获取引擎的JSON表示
    json_str = inspector.get_engine_information(trt.LayerInformationFormat.JSON)
    
    # 打印层信息
    print("引擎层信息:")
    print(json_str)
    
    return json_str

使用Tensor Cores

Tensor Cores是NVIDIA Volta、Turing和Ampere架构GPU中的专用硬件单元，可以显著加速矩阵乘法和卷积操作。TensorRT会自动使用Tensor Cores（如果可用）来加速计算。

要启用Tensor Cores，只需在构建引擎时启用FP16或INT8精度：

python 复制代码 
# 启用Tensor Cores
config.set_flag(trt.BuilderFlag.FP16)  # 对于FP16精度
# 或
config.set_flag(trt.BuilderFlag.INT8)  # 对于INT8精度

多流并行处理

在处理多个输入时，可以使用CUDA流来并行执行推理，提高吞吐量。

python 复制代码 
# 使用多流并行处理
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import threading

class MultiStreamInferencer:
    """使用多流的TensorRT推理器"""
    
    def __init__(self, engine_file_path, num_streams=4):
        """
        初始化推理器
        
        参数:
            engine_file_path: TensorRT引擎文件路径
            num_streams: 流的数量
        """
        # 创建TensorRT推理引擎
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        
        # 加载引擎
        with open(engine_file_path, 'rb') as f:
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        
        # 创建执行上下文
        self.contexts = [self.engine.create_execution_context() for _ in range(num_streams)]
        
        # 获取输入和输出的详细信息
        self.input_idx = self.engine.get_binding_index("input")
        self.output_idx = self.engine.get_binding_index("output")
        
        # 获取输入和输出形状
        self.input_shape = self.engine.get_binding_shape(self.input_idx)
        self.output_shape = self.engine.get_binding_shape(self.output_idx)
        
        # 创建CUDA流
        self.streams = [cuda.Stream() for _ in range(num_streams)]
        
        # 分配主机和设备内存
        self.host_inputs = []
        self.host_outputs = []
        self.device_inputs = []
        self.device_outputs = []
        self.bindings = []
        
        for _ in range(num_streams):
            host_input = cuda.pagelocked_empty(trt.volume(self.input_shape), dtype=np.float32)
            host_output = cuda.pagelocked_empty(trt.volume(self.output_shape), dtype=np.float32)
            
            device_input = cuda.mem_alloc(host_input.nbytes)
            device_output = cuda.mem_alloc(host_output.nbytes)
            
            bindings = [int(device_input), int(device_output)]
            
            self.host_inputs.append(host_input)
            self.host_outputs.append(host_output)
            self.device_inputs.append(device_input)
            self.device_outputs.append(device_output)
            self.bindings.append(bindings)
        
        # 创建线程锁
        self.lock = threading.Lock()
        
        # 创建线程池
        self.threads = []
        self.results = [None] * num_streams
    
    def infer_async(self, input_data_list):
        """
        异步执行推理
        
        参数:
            input_data_list: 输入数据列表
            
        返回:
            推理结果列表
        """
        num_inputs = len(input_data_list)
        num_streams = len(self.streams)
        
        # 确保输入数量不超过流的数量
        if num_inputs > num_streams:
            raise ValueError(f"输入数量 ({num_inputs}) 超过了流的数量 ({num_streams})")
        
        # 清空线程列表
        self.threads = []
        self.results = [None] * num_streams
        
        # 为每个输入创建一个线程
        for i in range(num_inputs):
            thread = threading.Thread(target=self._infer_single, args=(i, input_data_list[i]))
            self.threads.append(thread)
            thread.start()
        
        # 等待所有线程完成
        for thread in self.threads:
            thread.join()
        
        # 返回结果
        return self.results[:num_inputs]
    
    def _infer_single(self, stream_idx, input_data):
        """
        在单个流上执行推理
        
        参数:
            stream_idx: 流索引
            input_data: 输入数据
        """
        # 获取对应的资源
        host_input = self.host_inputs[stream_idx]
        host_output = self.host_outputs[stream_idx]
        device_input = self.device_inputs[stream_idx]
        device_output = self.device_outputs[stream_idx]
        bindings = self.bindings[stream_idx]
        stream = self.streams[stream_idx]
        context = self.contexts[stream_idx]
        
        # 复制输入数据到主机内存
        np.copyto(host_input, input_data.ravel())
        
        # 将输入数据从主机复制到设备
        cuda.memcpy_htod_async(device_input, host_input, stream)
        
        # 执行推理
        context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
        
        # 将输出数据从设备复制回主机
        cuda.memcpy_dtoh_async(host_output, device_output, stream)
        
        # 同步流
        stream.synchronize()
        
        # 保存结果
        with self.lock:
            self.results[stream_idx] = host_output.reshape(self.output_shape).copy()
    
    def __del__(self):
        """清理资源"""
        # 释放设备内存
        for device_input, device_output in zip(self.device_inputs, self.device_outputs):
            device_input.free()
            device_output.free()

内存优化

TensorRT提供了多种内存优化技术，包括权重共享、激活值重用和内存池。

python 复制代码 
# 配置内存优化
import tensorrt as trt

def configure_memory_optimization(config):
    """
    配置内存优化
    
    参数:
        config: TensorRT构建器配置
    """
    # 设置最大工作空间大小
    config.max_workspace_size = 1 << 30  # 1GB
    
    # 启用激活值重用
    config.set_flag(trt.BuilderFlag.REFIT)
    
    # 启用内存池
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

性能分析和调试

TensorRT提供了多种工具来分析和调试模型性能。

使用CUDA性能分析器

python 复制代码 
# 使用CUDA性能分析器
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import time

class Profiler(trt.IProfiler):
    """TensorRT性能分析器"""
    
    def __init__(self):
        trt.IProfiler.__init__(self)
        self.layers = []
        self.times = []
    
    def report_layer_time(self, layer_name, ms):
        """
        报告层的执行时间
        
        参数:
            layer_name: 层名称
            ms: 执行时间（毫秒）
        """
        self.layers.append(layer_name)
        self.times.append(ms)
    
    def print_report(self):
        """打印性能报告"""
        total_time = sum(self.times)
        
        print("=== 性能分析报告 ===")
        print(f"总层数: {len(self.layers)}")
        print(f"总执行时间: {total_time:.2f} ms")
        
        # 按执行时间排序
        sorted_indices = sorted(range(len(self.times)), key=lambda i: self.times[i], reverse=True)
        
        print("\n前10个最耗时的层:")
        for i in range(min(10, len(sorted_indices))):
            idx = sorted_indices[i]
            layer_name = self.layers[idx]
            time_ms = self.times[idx]
            percentage = time_ms / total_time * 100
            print(f"  {i+1}. {layer_name}: {time_ms:.2f} ms ({percentage:.2f}%)")
        
        print("\n总执行时间: {total_time:.2f} ms")

def profile_engine(engine_file_path, input_data):
    """
    分析TensorRT引擎性能
    
    参数:
        engine_file_path: TensorRT引擎文件路径
        input_data: 输入数据
    """
    # 创建logger和runtime
    logger = trt.Logger(trt.Logger.WARNING)
    runtime = trt.Runtime(logger)
    
    # 加载引擎
    with open(engine_file_path, 'rb') as f:
        engine = runtime.deserialize_cuda_engine(f.read())
    
    # 创建执行上下文
    context = engine.create_execution_context()
    
    # 获取输入和输出的详细信息
    input_idx = engine.get_binding_index("input")
    output_idx = engine.get_binding_index("output")
    
    # 获取输入和输出形状
    input_shape = engine.get_binding_shape(input_idx)
    output_shape = engine.get_binding_shape(output_idx)
    
    # 分配主机和设备内存
    host_input = cuda.pagelocked_empty(trt.volume(input_shape), dtype=np.float32)
    host_output = cuda.pagelocked_empty(trt.volume(output_shape), dtype=np.float32)
    
    device_input = cuda.mem_alloc(host_input.nbytes)
    device_output = cuda.mem_alloc(host_output.nbytes)
    
    # 准备绑定
    bindings = [int(device_input), int(device_output)]
    
    # 复制输入数据到主机内存
    np.copyto(host_input, input_data.ravel())
    
    # 将输入数据从主机复制到设备
    cuda.memcpy_htod(device_input, host_input)
    
    # 创建性能分析器
    profiler = Profiler()
    
    # 预热
    for _ in range(10):
        context.execute_v2(bindings=bindings)
    
    # 使用性能分析器执行推理
    context.execute_v2(bindings=bindings, profiler=profiler)
    
    # 将输出数据从设备复制回主机
    cuda.memcpy_dtoh(host_output, device_output)
    
    # 打印性能报告
    profiler.print_report()
    
    # 测量端到端延迟
    num_runs = 100
    start_time = time.time()
    
    for _ in range(num_runs):
        context.execute_v2(bindings=bindings)
    
    end_time = time.time()
    
    # 计算平均延迟
    avg_latency = (end_time - start_time) * 1000 / num_runs
    print(f"平均端到端延迟: {avg_latency:.2f} ms")
    
    # 释放设备内存
    device_input.free()
    device_output.free()

使用TensorRT的Inspector API

TensorRT 8.0引入了Inspector API，可以用来检查引擎的内部结构和配置：

python 复制代码 
# 使用TensorRT的Inspector API
def inspect_engine(engine_file_path):
    """
    检查TensorRT引擎
    
    参数:
        engine_file_path: TensorRT引擎文件路径
    """
    # 创建logger和runtime
    logger = trt.Logger(trt.Logger.WARNING)
    runtime = trt.Runtime(logger)
    
    # 加载引擎
    with open(engine_file_path, 'rb') as f:
        engine = runtime.deserialize_cuda_engine(f.read())
    
    # 创建检查器
    inspector = engine.create_engine_inspector()
    
    # 获取引擎信息
    engine_info = inspector.get_engine_information(trt.LayerInformationFormat.JSON)
    
    # 打印引擎信息
    print("=== 引擎信息 ===")
    print(engine_info)

最佳实践和优化技巧

以下是一些使用TensorRT进行模型优化和部署的最佳实践和技巧：

1. 选择合适的精度

根据应用需求和硬件支持选择合适的精度：

  • 对于需要高精度的应用，使用FP32
  • 对于大多数应用，FP16通常能提供良好的精度和性能平衡
  • 对于对精度要求不高的应用，INT8可以提供最佳性能

2. 使用动态批处理

如果应用需要处理不同大小的批次，使用动态批处理可以避免为每个批次大小构建单独的引擎：

python 复制代码 
# 设置动态批处理
profile = builder.create_optimization_profile()
profile.set_shape(input_name, (1, *input_shape[1:]), (opt_batch, *input_shape[1:]), (max_batch, *input_shape[1:]))
config.add_optimization_profile(profile)

3. 优化输入数据处理

  • 使用页锁定内存（pinned memory）减少主机到设备的数据传输时间
  • 使用异步数据传输和CUDA流重叠计算和数据传输
  • 批量处理数据以提高吞吐量

4. 模型优化

  • 移除不必要的操作和层
  • 合并可以合并的层（例如，卷积+批量归一化+ReLU）
  • 使用量化感知训练提高量化模型的精度

5. 使用TensorRT插件

对于TensorRT原生不支持的操作，可以使用自定义插件：

python 复制代码 
# 定义自定义插件
class MyPlugin(trt.IPluginV2):
    # 实现插件接口
    pass

# 注册插件
plugin_registry = trt.get_plugin_registry()
plugin_creator = plugin_registry.get_plugin_creator("MyPlugin", "1")
plugin = plugin_creator.create_plugin("my_plugin_instance", trt.PluginFieldCollection([]))

# 在网络中使用插件
layer = network.add_plugin_v2([input_tensor], plugin)

6. 使用ONNX-TensorRT解析器

使用ONNX作为中间表示可以简化从各种深度学习框架到TensorRT的转换过程：

python 复制代码 
# 从ONNX模型构建TensorRT引擎
parser = trt.OnnxParser(network, logger)
with open(onnx_file_path, 'rb') as model:
    parser.parse(model.read())

7. 使用TF-TRT或PyTorch-TensorRT

对于TensorFlow或PyTorch用户，可以使用TF-TRT或PyTorch-TensorRT直接在原始框架中集成TensorRT优化：

python 复制代码 
# TF-TRT示例
import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt

# 转换TensorFlow模型为TF-TRT模型
converter = trt.TrtGraphConverterV2(input_saved_model_dir="saved_model")
converter.convert()
converter.save("trt_saved_model")

# 加载并使用TF-TRT模型
trt_model = tf.saved_model.load("trt_saved_model")
output = trt_model(input_data)
python 复制代码 
# PyTorch-TensorRT示例
import torch
import torch_tensorrt

# 转换PyTorch模型为TorchScript
scripted_model = torch.jit.script(model)

# 使用PyTorch-TensorRT编译模型
trt_model = torch_tensorrt.compile(scripted_model,
                                  inputs=[torch_tensorrt.Input((1, 3, 224, 224))],
                                  enabled_precisions={torch.float16})

# 使用编译后的模型
output = trt_model(input_data)

总结

本章深入探讨了TensorRT的高级特性和优化技巧，包括量化技术、动态形状支持、层间融合、张量核心加速、多流并行处理、内存优化、性能分析和调试等。通过应用这些技术和最佳实践，可以显著提高深度学习模型的推理性能，满足各种实际应用场景的需求。

在实际应用中，应根据具体需求和硬件条件选择合适的优化策略，并通过性能分析工具不断调整和优化，以达到最佳的性能和精度平衡。

相关推荐
技能咖20 分钟前
2025春招市场迎AI热潮：生成式人工智能（GAI）认证如何重构人才竞争力
人工智能
2301_764441331 小时前
基于BERT的序列到序列（Seq2Seq）模型，生成文本摘要或标题
人工智能·python·深度学习·bert
说私域1 小时前
开源链动2+1模式与AI智能名片赋能的S2B2C共享经济新生态
人工智能·微信·小程序·开源
蹦蹦跳跳真可爱5892 小时前
Python----计算机视觉处理（Opencv:霍夫变换)
人工智能·python·opencv·计算机视觉
菜鸡中的奋斗鸡→挣扎鸡2 小时前
c++ count方法
开发语言·c++
SunkingYang2 小时前
C++中将记录集的数据复制到Excel工作表中的CRange类CopyFromRecordset函数异常怎么捕获
c++·excel·方法·异常捕获·copyfrom·recordset
Angel Q.2 小时前
3D点云的深度学习网络分类（按照作用分类）
深度学习·3d·分类
livefan2 小时前
英伟达「虚拟轨道+AI调度」专利：开启自动驾驶3.0时代的隐形革命
人工智能·机器学习·自动驾驶
热门推荐
01RAG 实践- Ollama+RagFlow 部署本地知识库02如何在WPS和Word/Excel中直接使用DeepSeek功能03DeepSeek各版本说明与优缺点分析04本地化部署AI知识库：基于Ollama+DeepSeek+AnythingLLM保姆级教程05本地部署DeepSeek教程（Mac版本）06从零安装 LLaMA-Factory 微调 Qwen 大模型成功及所有的坑07苍穹外卖面试总结08FreeCAD傻瓜教程之创建参数化几何图形-螺旋体、平面、球体、椭球体、圆柱体、圆锥体、棱柱、椭圆09DeepSeek R1本地化部署 Ollama + Chatbox 打造最强 AI 工具10JS闭包不理解？一杯咖啡的时间带你轻松理解闭包！