从零到一,激活GPU的力量:使用TensorRT优化和执行深度学习模型,你的TensorRT入门指南

TensorRT学习笔记

TensorRT模型优化

模型优化与流程介绍

TensorRT 是由 NVIDIA 提供的一种高性能深度学习推理(Inference)优化器和运行时库,可用于加速深度学习模型。TensorRT 的主要目标是优化并行计算的使用,以充分利用 NVIDIA GPU 的计算能力。它适用于各种类型的深度学习模型,包括 CNN(卷积神经网络)、RNN(循环神经网络)等。

TensorRT 实现模型优化的方法和策略包括:

  1. 层和张量融合:通过融合网络的连续层,TensorRT 可以减少数据从 GPU 内存中的读写次数,从而提高性能。
  2. 精度调整:TensorRT 支持混合精度计算,可以选择 FP32、FP16 或者 INT8 来进行计算,以权衡精度和性能。
  3. 动态输入尺寸:TensorRT 支持动态输入尺寸,这意味着你可以用不同尺寸的输入在同一模型上进行推理。
  4. 内核自动调度:TensorRT 会自动选择最优的 CUDA 内核进行计算,这使得在不同的硬件上都能获得最好的性能。

大致流程一般包含以下步骤:

  1. 创建并配置 builder:使用 nvinfer1::createInferBuilder() 创建一个 builder 对象,用于之后的模型构建。
  2. 创建并配置 network:使用 nvinfer1::createNetworkV2() 创建一个 network 对象,用于描述模型。
  3. 设置模型输入尺寸:根据你的需求设置模型输入的可能范围,这有助于 TensorRT 对模型进行优化。
  4. 创建并配置 config:使用 nvinfer1::createBuilderConfig() 创建一个 config 对象,用于保存模型优化过程中的各种设置,如精度、最大 batch size 等。
  5. 创建 CUDA 流并设置 config 的 profile stream:使用 makeCudaStream() 创建 CUDA 流,然后使用 setProfileStream() 将其设置到 config 中,用于异步执行 GPU 操作。
  6. 构建并序列化模型:使用 buildSerializedNetwork() 创建并优化模型,然后将其序列化,用于之后的推理过程。

模型优化过程步骤编码

  1. 创建并配置 builder:首先创建了一个 IBuilder,然后使用这个 IBuilder 来创建一个 INetworkDefinition
// ========== 1. 创建builder:创建优化的执行引擎(ICudaEngine)的关键工具 ==========
// 在几乎所有使用TensorRT的场合都会使用到IBuilder
// 只要TensorRT来进行优化和部署,都需要先创建和使用IBuilder。
std::unique_ptr<nvinfer1::IBuilder> builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
if (!builder)
{
    
    
    std::cerr << "Failed to create build" << std::endl;
    return -1;
} 
std::cout << "Successfully to create builder!!" << std::endl;

IBuilder 是 TensorRT 中的一个接口,负责创建一个优化的 ICudaEngineICudaEngine 是一个优化后的可执行网络,可以用于实际执行推理任务。几乎在所有使用TensorRT的场合都会使用到IBuilder,因为它是创建优化的执行引擎(ICudaEngine)的关键工具。不论你的模型是什么,只要你想用TensorRT来进行优化和部署,都需要创建和使用IBuilder。因此,创建 IBuilder 是必要的步骤,因为它是创建 ICudaEngine 的关键工具。这里使用了 std::unique_ptr 来管理 IBuilder 的内存。

在创建时需要提供一个日志对象以接收错误、警告和其他信息。这个步骤是整个模型创建流程的起点,没有 builder,我们就无法创建网络(模型)。

  1. 创建并配置 network
// ========== 2. 创建network:builder--->network ==========
// 设置batch, 数据输入的批次量大小
// 显性设置batch
const unsigned int explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
std::unique_ptr<nvinfer1::INetworkDefinition> network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));
if (!network)
{
    
    
    std::cout << "Failed to create network" << std::endl;
    return -1;
}

创建 network 是因为我们需要一个 INetworkDefinition 对象来描述我们的模型。我们使用 buildercreateNetworkV2 方法来创建一个网络,并为其提供一组标志来指定网络的行为。在这个例子中,我们设置了 kEXPLICIT_BATCH 标志,以启用显式批处理。

Batch通常在你需要进行大量数据的推理时会用到。在训练和推理阶段,我们通常不会一次只处理一个样本,而是会一次处理一批样本,这样可以更有效地利用计算资源,特别是在GPU上进行计算时。同时,选择合适的Batch大小也是一个需要平衡的问题,一方面,较大的Batch大小可以更好地利用硬件并行性,提高计算效率,另一方面,较大的Batch大小会消耗更多的内存资源,所以需要根据具体情况进行选择。

  1. 设置模型输入尺寸

在 TensorRT 中,优化配置文件(Optimization Profile)用于描述模型的输入尺寸和动态尺寸范围。通过优化配置文件,可以告诉 TensorRT 输入数据的可能尺寸范围,使其可以创建一个适应各种输入尺寸的优化后的模型。

// 配置网络参数
// 我们需要告诉tensorrt我们最终运行时,输入图像的范围,batch size的范围。这样tensorrt才能对应为我们进行模型构建与优化。
nvinfer1::ITensor* input = network->getInput(0); // 获取了网络的第一个输入节点。
nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile(); // 创建了一个优化配置文件。
// 网络的输入节点就是模型的输入层,它接收模型的输入数据。
// 在 TensorRT 中,优化配置文件(Optimization Profile)用于描述模型的输入尺寸和动态尺寸范围。
// 通过优化配置文件,可以告诉 TensorRT 输入数据的可能尺寸范围,使其可以创建一个适应各种输入尺寸的优化后的模型。

// 设置最小尺寸
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, 3, 640, 640));
// 设置最优尺寸
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(1, 3, 640, 640));
// 设置最大尺寸
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(1, 3, 640, 640));

在许多神经网络模型中,可能会有多个输入节点。例如,在一个同时接受图像和元数据输入的模型中,可能有一个输入节点用于接收图像,另一个输入节点用于接收元数据。auto input = network->getInput(0);这行代码获取了网络的第一个输入节点。网络的输入节点就是模型的输入层,它接收模型的输入数据。这里的 0 是一个索引,指的是网络的第一个输入节点。getInput 函数的参数就是用来指定你想获取哪个输入节点。如果你传入 0,那么就会返回第一个输入节点。如果你传入 1,那么就会返回第二个输入节点,以此类推。所以,你可以根据你的网络结构和需要,传入不同的参数来获取不同的输入节点。在你给出的这段代码中,因为网络可能只有一个输入节点,所以传入 0 来获取这个输入节点。

在这个步骤中,我们首先获取模型的输入节点(在这个例子中,我们假设模型只有一个输入),然后创建一个优化配置文件(OptimizationProfile)。优化配置文件描述了模型输入的可能范围,这对于模型的优化过程是必要的。我们为每个维度设置最小、最优和最大尺寸。

在这段代码中,每个网络层都有一个唯一的名字,getName 方法用于获取该输入节点(或者更一般地说,该张量)的名字。这个名字在定义网络模型时被赋予,通常用来帮助标识和跟踪网络中的各个节点。这个名字在 TensorRT 中有重要的用途,因为在设置输入输出节点的尺寸,或者在执行推理时,都会使用到这个名字。

  1. 创建并配置 config
// ========== 3. 创建config配置:builder--->config ==========
// 配置解析器
std::unique_ptr<nvinfer1::IBuilderConfig> config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
if (!config)
{
    
    
    std::cout << "Failed to create config" << std::endl;
    return -1;
}
// 添加之前创建的优化配置文件(profile)到配置对象(config)中
// 优化配置文件(profile)包含了输入节点尺寸的设置,这些设置会在模型优化时被使用。
config->addOptimizationProfile(profile);
// 设置精度
config->setFlag(nvinfer1::BuilderFlag::kFP16);
builder->setMaxBatchSize(1);

在这个步骤中,我们创建一个 IBuilderConfig 对象,用于保存模型优化过程中的各种设置。我们向其添加我们之前创建的优化配置文件,并设置一些标志,如 kFP16,来启用对 FP16 精度的支持。我们还设置了最大批处理大小,这是优化过程的一部分。

在 TensorRT 中,配置文件(IBuilderConfig 对象)用于保存模型优化过程中的各种设置。这些设置包括:

  • 优化精度模式:你可以选择 FP32、FP16 或 INT8 精度模式。使用更低的精度(如 FP16 或 INT8)可以减少计算资源的使用,从而提高推理速度,但可能会牺牲一些推理精度。
  • 最大批处理大小:这是模型优化过程中可以处理的最大输入批处理数量。
  • 工作空间大小:这是 TensorRT 在优化和执行模型时可以使用的最大内存数量。
  • 优化配置文件:优化配置文件描述了模型输入尺寸的可能范围。根据这些信息,TensorRT 可以创建一个针对各种输入尺寸优化的模型。

配置文件还可以包含其他的设置,例如 GPU 设备选择、层策略选择等等。这些设置都会影响 TensorRT 如何优化模型,以及优化后的模型的性能。

  1. 创建 CUDA 流并设置 config 的 profile stream
// 创建流,用于设置profile
auto profileStream = samplesCommon::makeCudaStream();
config->setProfileStream(*profileStream);

这里我们创建一个 CUDA 流,用于异步执行 GPU 操作。然后我们把这个 CUDA 流设置为配置对象的 profile stream。Profile stream 用于确定何时可以开始收集内核时间,以及何时可以读取收集到的时间。

这里的 profileStream 是 CUDA 流(CUDA Stream),而不是前面的优化配置文件(optimization profile)。CUDA 流是一个有序的操作队列,它们在 GPU 上异步执行。流可以用于组织和控制执行的并发性和依赖关系。

在你的代码中,profileStream 是通过 samplesCommon::makeCudaStream() 创建的。这个流会被传递给配置对象 config,然后 TensorRT 会在这个流中执行和优化配置文件(profile)相关的操作。这允许这些操作并发执行,从而提高了执行效率。

需要注意的是,这里的 CUDA 流和前面的优化配置文件是两个完全不同的概念。优化配置文件包含了模型输入尺寸的信息,用于指导模型的优化过程;而 CUDA 流则是用于管理 GPU 操作的执行顺序,以提高执行效率。两者虽然名字类似,但实际上在功能和用途上是完全不同的。

  1. 构建并序列化模型
// ========== 5. 序列化保存engine ==========
// 使用之前创建并配置的 builder、network 和 config 对象来构建并序列化一个优化过的模型。
auto plan = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));
std::ofstream engine_file("./weights/yolov5.engine", std::ios::binary);
engine_file.write((char *)plan->data(), plan->size());
engine_file.close();

在这最后一步中,我们使用 buildernetworkconfig 来构建并序列化我们的模型。这个过程可能会对模型进行一系列的优化,包括层融合、卷积算法选择、张量格式选择等。这个步骤产生了一个序列化的模型,我们可以将它保存到文件中,然后在之后的推理过程中加载和使用。这可以大大减少模型的加载时间,因为我们不需要再次进行优化过程。

使用之前创建并配置的 buildernetworkconfig 对象来构建并序列化一个优化过的模型。buildSerializedNetwork 函数返回的是一个 IHostMemory 对象,这个对象包含了序列化后的模型数据。

auto plan = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));

这行代码中,plan 是一个独特的指针,指向了序列化模型数据的内存。

这段代码的主要目的是创建一个序列化的模型,并将其保存到文件中。将模型序列化并保存到文件的好处是,你可以在程序的后续运行中直接加载这个文件,而不需要每次都重新构建和优化模型,这样可以大大提高效率。

知识点笔记

  • IBuilder 是用于创建 INetworkDefinition 的对象,也就是用于创建模型的对象。
  • INetworkDefinition 表示了一个 TensorRT 网络,用于描述模型的结构。
  • 显性批处理意味着在构建模型时显式指定批处理大小,而不是在运行时确定。这在某些情况下可以提高模型的性能。
  • IBuilderConfig 用于保存模型优化过程中的各种设置,包括优化精度模式、最大批处理大小、工作空间大小和优化配置文件等。
  • 优化配置文件描述了模型输入尺寸的可能范围,可以指导模型的优化过程。
  • CUDA 流是一个有序的操作队列,它们在 GPU 上异步执行。流可以用于组织和控制执行的并发性和依赖关系。
  • 使用 buildSerializedNetwork 方法可以创建并序列化一个模型,返回的是一个 IHostMemory 对象,这个对象包含了序列化后的模型数据。将模型序列化并保存到文件可以提高模型的加载速度。
  • 在 TensorRT 中,所有的模型输入都是网络层,每个网络层都有一个唯一的名字,可以使用 getName 方法获取。
TensorRT模型优化
用于创建 INetworkDefinition 对象的接口
用于描述模型的对象
设置模型输入的可能范围,用于模型优化
保存模型优化过程中的各种设置
用于异步执行 GPU 操作,确定何时可以开始收集内核时间,何时可以读取收集到的时间
进行模型优化,生成序列化模型,用于之后的推理过程
开始
创建并配置 builder
创建并配置 network
设置模型输入尺寸
创建并配置 config
创建 CUDA 流并设置 config 的 profile stream
构建并序列化模型
结束

完整代码

build.cu

#include <iostream>
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "logger.h"
#include "common.h"
#include "buffers.h"
#include <cassert>
#include <memory>

int main(int argc, char** argv)
{
    
    
    if (argc != 2)
    {
    
    
        std::cerr << "请输入onnx文件位置: ./build/[onnx_file]" << std::endl;
        return -1;
    }
    char* onnx_file = argv[1];
    // ========== 1. 创建builder:创建优化的执行引擎(ICudaEngine)的关键工具 ==========
    // 在几乎所有使用TensorRT的场合都会使用到IBuilder
    // 只要TensorRT来进行优化和部署,都需要先创建和使用IBuilder。
    auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
    if (!builder)
    {
    
    
        std::cerr << "Failed to create build" << std::endl;
        return -1;
    } 
    std::cout << "Successfully to create builder!!" << std::endl;

    // ========== 2. 创建network:builder--->network ==========
    // 设置batch, 数据输入的批次量大小
    // 显性设置batch
    const unsigned int explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    std::unique_ptr<nvinfer1::INetworkDefinition> network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));
    if (!network)
    {
    
    
        std::cout << "Failed to create network" << std::endl;
        return -1;
    }

    // 创建onnxparser,用于解析onnx文件
    auto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
    // 调用onnxparser的parseFromFile方法解析onnx文件
    bool parsed = parser->parseFromFile(onnx_file, static_cast<int>(sample::gLogger.getReportableSeverity()));
    if (!parsed)
    {
    
    
        std::cerr << "Failed to parse onnx file!!" << std::endl;
        return -1;
    }
    // 配置网络参数
    // 我们需要告诉tensorrt我们最终运行时,输入图像的范围,batch size的范围。这样tensorrt才能对应为我们进行模型构建与优化。
    nvinfer1::ITensor* input = network->getInput(0); // 获取了网络的第一个输入节点。
    nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile(); // 创建了一个优化配置文件。
    // 网络的输入节点就是模型的输入层,它接收模型的输入数据。
    // 在 TensorRT 中,优化配置文件(Optimization Profile)用于描述模型的输入尺寸和动态尺寸范围。
    // 通过优化配置文件,可以告诉 TensorRT 输入数据的可能尺寸范围,使其可以创建一个适应各种输入尺寸的优化后的模型。

    // 设置最小尺寸
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, 3, 640, 640));
    // 设置最优尺寸
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(1, 3, 640, 640));
    // 设置最大尺寸
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(1, 3, 640, 640));

    // ========== 3. 创建config配置:builder--->config ==========
    // 配置解析器
    std::unique_ptr<nvinfer1::IBuilderConfig> config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    if (!config)
    {
    
    
        std::cout << "Failed to create config" << std::endl;
        return -1;
    }
    // 添加之前创建的优化配置文件(profile)到配置对象(config)中
    // 优化配置文件(profile)包含了输入节点尺寸的设置,这些设置会在模型优化时被使用。
    config->addOptimizationProfile(profile);
    // 设置精度
    config->setFlag(nvinfer1::BuilderFlag::kFP16);
    builder->setMaxBatchSize(1);
    config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 30);

    // 创建流,用于设置profile
    auto profileStream = samplesCommon::makeCudaStream();
    if (!profileStream)
    {
    
    
        std::cerr << "Failed to create CUDA profileStream File" << std::endl;
        return -1;
    }
    config->setProfileStream(*profileStream);

    // ========== 5. 序列化保存engine ==========
    // 使用之前创建并配置的 builder、network 和 config 对象来构建并序列化一个优化过的模型。
    std::unique_ptr<nvinfer1::IHostMemory> plan = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));
    std::ofstream engine_file("./weights/best.engine", std::ios::binary);
    assert(engine_file.is_open() && "Failed to open engine file");
    engine_file.write((char *)plan->data(), plan->size());
    engine_file.close();

    // ========== 6. 释放资源 ==========
    std::cout << "Engine build success!" << std::endl;
    return 0;
}

TensorRT模型推理

模型部署与推理过程步骤编码

  1. 创建推理运行时的runtimeIRuntime 是 TensorRT 提供的一个接口,主要用于在推理阶段执行序列化的模型。创建 IRuntime 实例是在推理阶段加载和运行 TensorRT 引擎的首要步骤。如果你需要创建一个 IRuntime 实例。createInferRuntime 函数是 TensorRT 提供的一个函数,用于创建 IRuntime 实例。这个函数需要一个 ILogger 实例作为参数,ILogger 主要用于处理 TensorRT 在运行过程中的日志信息。
// 使用 std::unique_ptr 是为了管理 IRuntime 实例的生命周期。
auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(sample::gLogger.getTRTLogger()));
if (!runtime)	// 检查IRuntime 实例创建是否成功。
{
    std::cerr << "Failed to create runtime." << std::endl;
    return -1;
}
在使用TensorRT进行推理时,你需要用到`IRuntime`。在TensorRT的工作流程中,一般有两个主要阶段:优化阶段和推理阶段。
  • 在优化阶段,我们使用IBuilderINetworkDefinition等接口创建和优化模型。优化完成后,我们通常将优化后的模型序列化并保存到硬盘上。
  • 在推理阶段,我们需要从硬盘上加载优化后的模型,然后执行推理。这个阶段就需要用到IRuntime。我们首先使用IRuntimedeserializeCudaEngine方法从序列化的数据中加载模型,然后使用加载的模型进行推理。

所以,只要需要使用TensorRT进行推理,你就需要使用IRuntime。这包括在应用程序中进行推理,或者在模型训练过程中进行推理以评估模型性能等场景。

  1. 读取模型文件并反序列生成engine与创建ICudaEngine神经网络模型对象:我们首先加载了保存在硬盘上的模型文件,然后使用IRuntimedeserializeCudaEngine方法将其反序列化为ICudaEngine对象。这是因为在TensorRT中,推理的执行是通过ICudaEngine对象来进行的。

读取模型文件

// 读取模型文件的函数
void load_engine_file(const char* engine_file, std::vector<uchar>& engine_data)
{
    // 初始化engine_data,'\0'表示空字符
    engine_data = { '\0' };
    // 打开模型文件,以二进制模式打开
    std::ifstream engine_fp(engine_file, std::ios::binary);
    if (!engine_fp.is_open())	// 如果文件未成功打开,输出错误信息并退出程序
    {
        std::cerr << "Unable to load engine file." << std::endl;
        exit(-1);
    }
    engine_fp.seekg(0, engine_fp.end); // 将文件指针移动到文件末尾,用于获取文件大小
    int length = engine_fp.tellg();	   // 获取文件大小
    engine_data.resize(length);		   // 根据文件大小调整engine_data的大小
    engine_fp.seekg(0, engine_fp.beg); // 将文件指针重新定位到文件开始位置
    // 读取文件内容到engine_data中,reinterpret_cast<char*>是用来将uchar*类型指针转换为char*类型指针
    engine_fp.read(reinterpret_cast<char*> (engine_data.data()), length);
    engine_fp.close();// 关闭文件
}

反序列读取后实例化为Engine对象

在TensorRT中,推理的执行是通过ICudaEngine对象来进行的。

// 加载了保存在硬盘上的模型文件
// 存储到std::vector<uchar>类型的engine_data变量中,以便于后续的模型反序列化操作。
std::vector<uchar> engine_data = { '\0' };
load_engine_file(engine_file, engine_data);
// 使用IRuntime的deserializeCudaEngine方法将其反序列化为ICudaEngine对象。
auto mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(engine_data.data(), engine_data.size()));
if (!mEngine)
{
    std::cerr << "Failed to create engine." << std::endl;
    return -1;
}

在上面的代码中,ICudaEngine是一个重要的接口,它封装了优化后的模型和运行推理所需的所有资源,包括后续的执行上下文、输入输出内存管理、执行推理的CUDA流等待。所以在执行推理前,我们需要先通过反序列化得到ICudaEngine对象。在这个过程中,我们将之前优化和序列化的模型文件加载到内存中,然后通过IRuntimedeserializeCudaEngine方法将其转化为ICudaEngine对象。

  1. 创建执行上下文:在 TensorRT 中,执行推理的过程是通过执行上下文 (ExecutionContext) 来进行的。

在TensorRT中,ICudaEngine对象代表了优化后的网络,而IExecutionContext则封装了执行推理所需要的上下文信息,比如输入/输出的内存、CUDA流等。每个IExecutionContext都与一个特定的ICudaEngine相关联,并且在执行推理时会使用ICudaEngine中的模型。

创建IExecutionContext的过程是通过调用ICudaEnginecreateExecutionContext()方法完成的。

// 通过调用ICudaEngine`的`createExecutionContext()方法创建对应的上下文管理器
auto context = std::unique_ptr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
if (!context)
{
    
    
    std::cerr << "Failed to create ExcutionContext." << std::endl;
    return -1;
}

在 TensorRT 中,推理的执行是通过执行上下文 (ExecutionContext) 来进行的.ExecutionContext封装了推理运行时所需要的所有信息,包括存储设备上的缓冲区地址、内核参数以及其他设备信息。因此,当你想在设备上执行模型推理时必须要有一个ExecutionContext

每一个ExecutionContext是和一个特定的ICudaEngine(优化后的网络) 相关联的,它包含了网络推理的所有运行时信息。因此,没有执行上下文,就无法进行模型推理。

此外,每个ExecutionContext也和一个特定的CUDA流关联,它允许在同一个流中并行地执行多个模型推理,这使得能够在多个设备或多个线程之间高效地切换。

深化理解ICudaEngine

ICudaEngine 是 NVIDIA TensorRT 库中的一个关键接口,它提供了在 GPU 上执行推断的所有必要的方法。每个 ICudaEngine 对象都表示了一个已经优化过的神经网络模型。

以下是一些你可以使用 ICudaEngine 完成的操作:

  1. 创建执行上下文:使用 ICudaEngine 对象,你可以创建一个或多个执行上下文 (IExecutionContext),每个上下文在给定的引擎中定义了网络的一次执行。这对于并发执行或使用不同的批量大小进行推断很有用。
  2. 获取网络层信息:ICudaEngine 提供了方法来查询网络的输入和输出张量以及网络层的相关信息,例如张量的尺寸、数据类型等。
  3. 序列化和反序列化:你可以将 ICudaEngine 对象序列化为一个字符串,然后将这个字符串保存到磁盘,以便以后使用。这对于保存优化后的模型并在以后的会话中重新加载它们很有用。相对应地,你也可以使用 deserializeCudaEngine() 方法从字符串或磁盘文件中恢复 ICudaEngine 对象。

通常,ICudaEngine 用于以下目的:

  1. 执行推断:最主要的应用就是使用 ICudaEngine 执行推断。你可以创建一个执行上下文,并使用它将输入数据提供给模型,然后从模型中获取输出结果。
  2. 加速模型加载:通过序列化 ICudaEngine 对象并将它们保存到磁盘,你可以在以后的会话中快速加载优化后的模型,无需重新进行优化。
  3. 管理资源:在并发执行或使用不同的批量大小进行推断时,你可以创建多个执行上下文以管理 GPU 资源。

总的来说,ICudaEngine 是 TensorRT 中最重要的接口之一,它提供了一种高效、灵活的方式来在 GPU 上执行推断。

  1. 创建输入输出缓冲区:输入输出缓冲区在深度学习推理中起着桥梁的作用,它们连接着原始输入数据和模型处理结果,使得我们可以有效地执行模型推理,并获取推理结果。

输入输出缓冲区在深度学习推理中发挥着至关重要的作用,它们充当了处理过程中数据的"桥梁",使得我们能够将输入数据顺利地送入模型,并从模型中获取处理结果。这大大简化了深度学习推理的过程,使得我们可以更专注于实现模型的逻辑,而不需要过多地关心数据的传输和存储问题。

TensorRT 提供的 BufferManager 类用于简化这个过程,可以自动创建并管理这些缓冲区,使得在进行推理时不需要手动创建和管理这些缓冲区。

TensorRT中的BufferManager类是一个辅助类,它的主要目的是简化输入/输出缓冲区的创建和管理。在进行模型推理时,不需要手动管理这些缓冲区,只需要将输入数据放入BufferManager管理的输入缓冲区中,然后调用推理函数。待推理完成后,可以从BufferManager管理的输出缓冲区中获取模型的推理结果。

// 4. 创建输入输出缓冲区
samplesCommon::BufferManager buffers(mEngine);

在上述代码中,samplesCommon::BufferManager buffers(mEngine);这行代码创建了一个BufferManager对象buffers,并用模型引擎mEngine来初始化它。这个BufferManager对象会为模型引擎的输入和输出创建对应的缓冲区,以便于进行后续的模型推理操作。

在深度学习推理过程中,输入输出缓冲区是必不可少的,其主要用途如下:

输入缓冲区: 输入缓冲区主要用于存储需要进行处理的数据。在深度学习推理中,输入缓冲区通常用于存储模型的输入数据。比如,如果你的模型是一个图像识别模型,那么输入缓冲区可能会存储待识别的图像数据。当执行模型推理时,模型会从输入缓冲区中读取数据进行处理。

输出缓冲区: 输出缓冲区主要用于存储处理过的数据,即处理结果。在深度学习推理中,输出缓冲区通常用于存储模型的输出结果。继续上述图像识别模型的例子,一旦模型完成了图像的识别处理,识别结果(例如,图像中物体的类别)就会存储在输出缓冲区中。我们可以从输出缓冲区中获取这些结果,进行进一步的处理或分析。

总的来说,输入输出缓冲区在深度学习推理中起着桥梁的作用,它们连接着原始输入数据和模型处理结果,使得我们可以有效地执行模型推理,并获取推理结果。

  1. **读入视频与申请GPU内存:**在前面4个步骤中,我们已经完成了基本的模型推理准备工作的代码编写了,现在我们要开始读入视频流并进行推理。同时,我们还需要申请对应的GPU内存空间,GPU 内存用于存储需要在 GPU 上进行处理的数据。在图像处理和深度学习等计算密集型任务中,GPU 的并行处理能力可以大大加快运算速度。因此,我们通常会将数据存储在 GPU 内存中,以便利用 GPU 的计算能力。
// 5.读入视频
auto cap = cv::VideoCapture(input_path_file);
int width = int(cap.get(cv::CAP_PROP_FRAME_WIDTH));
int height = int(cap.get(cv::CAP_PROP_FRAME_HEIGHT));
int fps = int(cap.get(cv::CAP_PROP_FPS));

// 写入MP4文件,参数分别是:文件名,编码格式,帧率,帧大小
cv::VideoWriter writer("./output/test.mp4", cv::VideoWriter::fourcc('H', '2', '6', '4'), fps, cv::Size(width, height));

cv::Mat frame;
int frame_index = 0;
// 申请gpu内存
cuda_preprocess_init(height * width);

上面的代码主要做了几件事:

  1. 打开一个视频文件并获取其属性:cv::VideoCapture对象被用来打开输入的视频文件。通过调用cv::CAP_PROP_FRAME_WIDTHcv::CAP_PROP_FRAME_HEIGHTcv::CAP_PROP_FPS方法,代码获取了视频的宽度、高度和帧率。
  2. 初始化视频写入对象:cv::VideoWriter对象被用来创建一个新的视频文件,编码格式为H264,帧率和帧大小与输入的视频相同。
  3. 定义一个cv::Mat对象来保存每一帧的图像数据,以及一个整型变量frame_index来记录当前处理的帧的索引。
  4. 申请GPU内存:cuda_preprocess_init(height * width)函数用于在GPU上申请一块内存空间,该内存大小等于视频帧的像素数。这是因为在后续的视频处理过程中,我们可能需要在GPU上对每一个像素进行操作,因此需要申请足够的GPU内存来存储这些像素数据。

什么是GPU内存?为什么需要申请?申请多大?

GPU 内存用于存储需要在 GPU 上进行处理的数据。在图像处理和深度学习等计算密集型任务中,GPU 的并行处理能力可以大大加快运算速度。因此,我们通常会将数据存储在 GPU 内存中,以便利用 GPU 的计算能力。

在上面的代码中,cuda_preprocess_init(height * width) 这一行代码用于在 GPU 上申请一段内存空间,这段内存的大小等于视频帧的像素数(宽度乘以高度)。这是因为,对于图像数据(包括视频帧),每一个像素通常都需要进行处理(比如色彩空间转换、缩放、归一化等),因此需要为每一个像素申请一份 GPU 内存。

选择申请多大的 GPU 内存取决于你的需求。如果你正在处理的是图像数据,那么通常你需要为每个像素申请一份内存。如果你正在处理的是其他类型的数据,那么需要根据数据的特点来决定。但是请注意,GPU 内存是有限的,过大的内存需求可能会超出 GPU 的容量。在实际应用中,需要根据实际情况和具体需求来合理选择申请的 GPU 内存大小。

对于一般的图像处理任务,一种常见的做法是为每个像素申请一份内存。对于深度学习任务,可能需要更多的内存来存储中间计算结果和模型的参数。

  1. 图像预处理、推理与结果输出:上面的过程我们已经完成了推理前的所有操作了,现在我们要开始正式的视频推理操作。这个过程包括预处理、模型推理、结果提取和后处理等步骤。

模型推理是深度学习的关键步骤,它将训练好的模型应用到新的数据上,进行预测。在本处例子中,我们将会进行如下步骤:

  1. **输入预处理:**将输入的图像数据转化为模型需要的格式。这包括letterbox(尺度变换)、归一化、BGR转RGB等操作。这是为了确保模型可以正确解读输入数据。
  2. **执行推理:**运行模型进行预测,这是利用训练好的模型进行预测的关键步骤。在这个例子中,模型是用来做目标检测的。
  3. **拷贝输出:**将推理的结果从设备上的缓冲区复制到主机上的缓冲区。这是为了能在主机上对推理结果进行进一步处理。
  4. **提取推理结果:**从主机上的缓冲区中获取推理结果,这包括了目标检测的数量、类别、置信度以及边界框等信息。这使得我们可以获取模型的预测结果。
  5. **非极大值抑制:**对预测的结果进行后处理,消除冗余的检测框,只保留最佳的检测结果。这是为了提高目标检测的准确性。
// 输入预处理(实现了对输入图像处理的gpu 加速)
process_input(frame, (float*)buffers.getDeviceBuffer(kInputTensorName));
// 6. 执行推理
context->executeV2(buffers.getDeviceBindings().data());
// 推理完成后拷贝回缓冲区
buffers.copyOutputToHost();
// 从buffer中提取推理结果
int32_t* num_det = (int32_t*)buffers.getHostBuffer(kOutNumDet); // 目标检测到的个数
int32_t* cls = (int32_t*)buffers.getHostBuffer(kOutDetCls);     // 目标检测到的目标类别
float* conf = (float*)buffers.getHostBuffer(kOutDetScores);     // 目标检测的目标置信度
float* bbox = (float*)buffers.getHostBuffer(kOutDetBBoxes);     // 目标检测到的目标框
// 非极大值抑制,得到最后的检测框
std::vector<Detection> bboxs;
yolo_nms(bboxs, num_det, cls, conf, bbox, kConfThresh, kNmsThresh);

这段代码主要用于执行深度学习模型的推理以及后处理。以下是如上代码步骤的解释:

  1. 输入预处理:process_input函数用于预处理输入帧,即将原始的图像数据转化为模型需要的输入格式。这一步通常包括缩放、裁剪、归一化等操作。这里使用了GPU加速,以加快预处理速度。处理后的数据存储在设备(GPU)上的缓冲区中,这个缓冲区通过buffers.getDeviceBuffer(kInputTensorName)获取。
  2. 执行推理:context->executeV2(buffers.getDeviceBindings().data())函数执行模型的推理,这一步是整个过程中最关键的一步。这里的context对象是已经被配置和优化好的TensorRT执行上下文,buffers.getDeviceBindings().data()返回的是输入输出缓冲区在设备(GPU)上的地址。
  3. 拷贝输出:buffers.copyOutputToHost()函数将推理的结果从设备(GPU)上的缓冲区复制到主机(CPU)上的缓冲区。
  4. 提取推理结果:使用buffers.getHostBuffer()函数从主机上的缓冲区中获取推理的结果。这里的结果包括了目标检测的数量、类别、置信度以及边界框等信息。
  5. 非极大值抑制:yolo_nms函数进行非极大值抑制,这是目标检测后处理的常见步骤,用于消除冗余的检测框,只保留最佳的检测结果。kConfThreshkNmsThresh是非极大值抑制的两个阈值,一般会通过实验来确定。

总的来说,这段代码通过执行预处理、推理、后处理等步骤,实现了对输入图像的目标检测任务,并生成了最终的检测结果。

在如上代码中,会有一些特别的变量,通常由TensorRT社区一起开发

// These are used to define input/output tensor names,
// you can set them to whatever you want.
const static char* kInputTensorName = "images";
const static char* kOutputTensorName = "prob";
const static char* kOutNumDet = "DecodeNumDetection";
const static char* kOutDetScores = "DecodeDetectionScores";
const static char* kOutDetBBoxes = "DecodeDetectionBoxes";
const static char* kOutDetCls = "DecodeDetectionClasses";

// Detection model and Segmentation model' number of classes
constexpr static int kNumClass = 80;

// Classfication model's number of classes
constexpr static int kClsNumClass = 1000;

constexpr static int kBatchSize = 1;

// Yolo's input width and height must by divisible by 32
constexpr static int kInputH = 640;
constexpr static int kInputW = 640;

// Classfication model's input shape
constexpr static int kClsInputH = 224;
constexpr static int kClsInputW = 224;

// Maximum number of output bounding boxes from yololayer plugin.
// That is maximum number of output bounding boxes before NMS.
constexpr static int kMaxNumOutputBbox = 1000;

constexpr static int kNumAnchor = 3;

// The bboxes whose confidence is lower than kIgnoreThresh will be ignored in yololayer plugin.
constexpr static float kIgnoreThresh = 0.1f;

/* --------------------------------------------------------
 * These configs are NOT related to tensorrt model, if these are changed,
 * please re-compile, but no need to re-serialize the tensorrt model.
 * --------------------------------------------------------*/

// NMS overlapping thresh and final detection confidence thresh
const static float kNmsThresh = 0.7f;
const static float kConfThresh = 0.4f;

const static int kGpuId = 0;

// If your image size is larger than 4096 * 3112, please increase this value
const static int kMaxInputImageSize = 4096 * 3112;

上面的变量中,通常用来表示:

  1. kInputTensorName:定义模型输入张量的名字。
  2. kOutputTensorName:定义模型输出张量的名字。
  3. kOutNumDet:定义了模型输出的部分名称,表示目标检测的数量。
  4. kOutDetScores:定义了模型输出的部分名称,表示每个检测到的目标的置信度。
  5. kOutDetBBoxes:定义了模型输出的部分名称,表示每个检测到的目标的边界框。
  6. kOutDetCls:定义了模型输出的部分名称,表示每个检测到的目标的类别。
  7. kNumClass:定义了目标检测模型的类别数量。
  8. kClsNumClass:定义了分类模型的类别数量。
  9. kBatchSize:定义了每次推理的批次大小。
  10. kInputHkInputW:定义了输入图像的高度和宽度。
  11. kClsInputHkClsInputW:定义了分类模型的输入图像的高度和宽度。
  12. kMaxNumOutputBbox:定义了在执行非极大值抑制前,最大输出边界框的数量。
  13. kNumAnchor:定义了锚点的数量。
  14. kIgnoreThresh:定义了在 YOLO 层插件中,置信度低于该阈值的边界框会被忽略。
  15. kNmsThresh:定义了非极大值抑制的阈值。
  16. kConfThresh:定义了最终检测置信度的阈值。
  17. kGpuId:定义了要使用的 GPU 的 ID。
  18. kMaxInputImageSize:定义了输入图像的最大大小,如果图像大小超过这个值,需要增大这个值。

推理结果封装与后处理的部分代码:

主要包含了一个描述数据框的数据结构Detection和两个函数iouyolo_nms。这些主要用于实现目标检测任务中的非极大值抑制(Non-maximum Suppression,NMS)策略。NMS是一种在目标检测中常用的后处理方法,它用于过滤掉冗余的和重叠度较高的预测边界框。

#include <algorithm>
#include <vector>
#include <map>

// 定义检测框数据结构,包含bbox,conf和class_id三个字段
struct alignas(float) Detection
{
    
    
    float bbox[4];  // xmin ymin xmax ymax
    float conf;     // 置信度,表示模型对该检测框内物体的确信程度
    float class_id; // 类别id,表示该检测框内物体的类别
};

// 计算两个矩形框的IoU(Intersection over Union,交并比),是一种衡量矩形框重叠程度的度量
float iou(float lbox[4], float rbox[4])
{
    
    
    // 计算两个矩形框的交集
    float interBox[] = {
    
    
        (std::max)(lbox[0] - lbox[2] / 2.f, rbox[0] - rbox[2] / 2.f), // left
        (std::min)(lbox[0] + lbox[2] / 2.f, rbox[0] + rbox[2] / 2.f), // right
        (std::max)(lbox[1] - lbox[3] / 2.f, rbox[1] - rbox[3] / 2.f), // top
        (std::min)(lbox[1] + lbox[3] / 2.f, rbox[1] + rbox[3] / 2.f), // bottom
    };

    // 如果交集为空(即两个矩形框无交叠),则返回0
    if (interBox[2] > interBox[3] || interBox[0] > interBox[1])
    {
    
    
        return 0.0f;
    }

    // 计算交集面积
    float interBoxS = (interBox[1] - interBox[0]) * (interBox[3] - interBox[2]);

    // 计算IoU,即交集面积与并集面积的比值
    return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS);
}

// 实现非极大值抑制(Non-maximum Suppression,NMS),用于过滤掉冗余的和重叠度较高的预测边界框
void yolo_nms(std::vector<Detection> &res, int32_t *num_det, int32_t *cls, float *conf, float *bbox, float conf_thresh, float nms_thresh)
{
    
    
    // 创建一个空的map,用于存储每个类别的检测框
    std::map<float, std::vector<Detection>> m;
    for (int i = 0; i < num_det[0]; i++)
    {
    
    
        // 跳过置信度低于阈值的检测框
        if (conf[i] <= conf_thresh)
            continue;

        // 创建一个新的检测框,并将其添加到对应类别的检测框列表中
        Detection det;
        det.bbox[0] = bbox[i * 4 + 0];
        det.bbox[1] = bbox[i * 4 + 1];
        det.bbox[2] = bbox[i * 4 + 2];
        det.bbox[3] = bbox[i * 4 + 3];
        det.conf = conf[i];
        det.class_id = cls[i];
        if (m.count(det.class_id) == 0)
            m.emplace(det.class_id, std::vector<Detection>());
        m[det.class_id].push_back(det);
    }

    // 对每个类别的检测框列表进行处理
    for (auto it = m.begin(); it != m.end(); it++)
    {
    
    
        // 获取当前类别的检测框列表
        auto &dets = it->second;
        
        // 根据置信度对检测框列表进行排序
        std::sort(dets.begin(), dets.end(), cmp);
        
        // 对排序后的检测框列表进行处理
        for (size_t m = 0; m < dets.size(); ++m)
        {
    
    
            auto &item = dets[m];
            res.push_back(item);
            for (size_t n = m + 1; n < dets.size(); ++n)
            {
    
    
                // 如果两个检测框的IoU大于阈值,则删除后一个检测框
                if (iou(item.bbox, dets[n].bbox) > nms_thresh)
                {
    
    
                    dets.erase(dets.begin() + n);
                    --n;
                }
            }
        }
    }
}

IoU(Intersection over Union,交并比) 是一个衡量两个边界框重叠程度的度量。这个度量定义为两个边界框的交集面积除以它们的并集面积。IoU的值在0和1之间,值越大表示两个边界框的重叠程度越高。

非极大值抑制(Non-Maximum Suppression,NMS) 是目标检测中一种常用的后处理方法。在目标检测中,我们的模型可能会对同一个物体输出多个边界框,这些边界框可能有不同的位置和大小,但它们都预测到了同一个物体。因此,我们需要一种方法来从这些边界框中选择出最“好”的一个,这就是NMS要做的事情。

非极大值抑制的基本思想是:在多个重叠的边界框中,保留置信度最高的边界框,然后移除与它有高度重叠(通常使用IoU作为衡量重叠程度的指标)并且置信度较低的边界框。这个过程一直进行,直到所有的边界框都被检查过。最终,我们得到的边界框都是互不重叠的,并且每个物体都只被一个边界框预测到。

所以,IoU在非极大值抑制中用来判断两个边界框是否重叠,并且衡量它们的重叠程度。而非极大值抑制则用于从多个预测边界框中选择出最好的边界框。

操作流程图如下

createInferRuntime
Runtime创建成功
deserializeCudaEngine
Engine创建成功
createExecutionContext
Context创建成功
BufferManager类
缓冲区创建成功
VideoCapture类
cuda_preprocess_init
视频读取成功
GPU内存申请成功
process_input
executeV2
copyOutputToHost
getHostBuffer
yolo_nms
创建推理运行时的runtime
读取模型文件并反序列生成engine
创建执行上下文
创建输入输出缓冲区
读入视频与申请GPU内存
图像预处理 推理与结果输出
IRuntime 实例
优化阶段和推理阶段
ICudaEngine 实例
推理的执行
IExecutionContext 实例
执行推理
输入/输出缓冲区
执行模型推理, 并获取推理结果
读入视频流
申请GPU内存
准备视频推理
利用GPU进行处理
图像预处理
模型推理
结果提取
从buffer中提取推理结果
非极大值抑制,得到最后的检测框

完整代码

runtime.cu

#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "logger.h"
#include "common.h"
#include "buffers.h"
#include "utils/preprocess.h"
#include "utils/postprocess.h"
#include "utils/types.h"
#include <algorithm>
#include <vector>
#include <map>
static uint8_t* img_buffer_device = nullptr;
// 定义检测框数据结构,包含bbox,conf和class_id三个字段
#include <cuda_runtime_api.h>

#ifndef CUDA_CHECK
#define CUDA_CHECK(callstr)\
    {
      
      \
        cudaError_t error_code = callstr;\
        if (error_code != cudaSuccess) {
      
      \
            std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__;\
            assert(0);\
        }\
    }
#endif  // CUDA_CHECK

void cuda_preprocess_init(int max_image_size)
{
    
    
  // prepare input data in device memory
  CUDA_CHECK(cudaMalloc((void**)&img_buffer_device, max_image_size * 3));
}

void cuda_preprocess_destroy()
{
    
    
  CUDA_CHECK(cudaFree(img_buffer_device));
}
struct alignas(float) Detection
{
    
    
    float bbox[4];  // xmin ymin xmax ymax
    float conf;     // 置信度,表示模型对该检测框内物体的确信程度
    float class_id; // 类别id,表示该检测框内物体的类别
};

// 计算两个矩形框的IoU(Intersection over Union,交并比),是一种衡量矩形框重叠程度的度量
float iou(float lbox[4], float rbox[4])
{
    
    
    // 计算两个矩形框的交集
    float interBox[] = {
    
    
        (std::max)(lbox[0] - lbox[2] / 2.f, rbox[0] - rbox[2] / 2.f), // left
        (std::min)(lbox[0] + lbox[2] / 2.f, rbox[0] + rbox[2] / 2.f), // right
        (std::max)(lbox[1] - lbox[3] / 2.f, rbox[1] - rbox[3] / 2.f), // top
        (std::min)(lbox[1] + lbox[3] / 2.f, rbox[1] + rbox[3] / 2.f), // bottom
    };

    // 如果交集为空(即两个矩形框无交叠),则返回0
    if (interBox[2] > interBox[3] || interBox[0] > interBox[1])
    {
    
    
        return 0.0f;
    }

    // 计算交集面积
    float interBoxS = (interBox[1] - interBox[0]) * (interBox[3] - interBox[2]);

    // 计算IoU,即交集面积与并集面积的比值
    return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS);
}

// 实现非极大值抑制(Non-maximum Suppression,NMS),用于过滤掉冗余的和重叠度较高的预测边界框
void yolo_nms(std::vector<Detection> &res, int32_t *num_det, int32_t *cls, float *conf, float *bbox, float conf_thresh, float nms_thresh)
{
    
    
    // 创建一个空的map,用于存储每个类别的检测框
    std::map<float, std::vector<Detection>> m;
    for (int i = 0; i < num_det[0]; i++)
    {
    
    
        // 跳过置信度低于阈值的检测框
        if (conf[i] <= conf_thresh)
            continue;

        // 创建一个新的检测框,并将其添加到对应类别的检测框列表中
        Detection det;
        det.bbox[0] = bbox[i * 4 + 0];
        det.bbox[1] = bbox[i * 4 + 1];
        det.bbox[2] = bbox[i * 4 + 2];
        det.bbox[3] = bbox[i * 4 + 3];
        det.conf = conf[i];
        det.class_id = cls[i];
        if (m.count(det.class_id) == 0)
            m.emplace(det.class_id, std::vector<Detection>());
        m[det.class_id].push_back(det);
    }

    // 对每个类别的检测框列表进行处理
    for (auto it = m.begin(); it != m.end(); it++)
    {
    
    
        // 获取当前类别的检测框列表
        auto &dets = it->second;
        
        // 根据置信度对检测框列表进行排序
        std::sort(dets.begin(), dets.end(), cmp);
        
        // 对排序后的检测框列表进行处理
        for (size_t m = 0; m < dets.size(); ++m)
        {
    
    
            auto &item = dets[m];
            res.push_back(item);
            for (size_t n = m + 1; n < dets.size(); ++n)
            {
    
    
                // 如果两个检测框的IoU大于阈值,则删除后一个检测框
                if (iou(item.bbox, dets[n].bbox) > nms_thresh)
                {
    
    
                    dets.erase(dets.begin() + n);
                    --n;
                }
            }
        }
    }
}
// 读取模型文件的函数
void load_engine_file(const char* engine_file, std::vector<uchar>& engine_data)
{
    
    
    // 初始化engine_data,'\0'表示空字符
    engine_data = {
    
     '\0' };
    // 打开模型文件,以二进制模式打开
    std::ifstream engine_fp(engine_file, std::ios::binary);
    if (!engine_fp.is_open())	// 如果文件未成功打开,输出错误信息并退出程序
    {
    
    
        std::cerr << "Unable to load engine file." << std::endl;
        exit(-1);
    }
    engine_fp.seekg(0, engine_fp.end); // 将文件指针移动到文件末尾,用于获取文件大小
    int length = engine_fp.tellg();	   // 获取文件大小
    engine_data.resize(length);		   // 根据文件大小调整engine_data的大小
    engine_fp.seekg(0, engine_fp.beg); // 将文件指针重新定位到文件开始位置
    // 读取文件内容到engine_data中,reinterpret_cast<char*>是用来将uchar*类型指针转换为char*类型指针
    engine_fp.read(reinterpret_cast<char*> (engine_data.data()), length);
    engine_fp.close();// 关闭文件
}


int main(int argc, char** argv)
{
    
    
    if (argc < 3)
    {
    
    
        std::cerr << "需要2个参数, 请输入足够的参数, 用法: <engine_file> <input_path_file>" << std::endl;
        return -1;
    }
    // 在推理阶段,我们需要从硬盘上加载优化后的模型,然后执行推理。这个阶段就需要用到IRuntime。
    // 我们首先使用IRuntime的deserializeCudaEngine方法从序列化的数据中加载模型,然后使用加载的模型进行推理。
    const char* engine_file = argv[1];
    const char* input_path_file = argv[2];

    // 1. 创建推理运行时的runtime
    // IRuntime 是 TensorRT 提供的一个接口,主要用于在推理阶段执行序列化的模型。
    // 创建 IRuntime 实例是在推理阶段加载和运行 TensorRT 引擎的首要步骤。
    auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(sample::gLogger.getTRTLogger()));
    if (!runtime)
    {
    
    
        std::cerr << "Failed to create runtime." << std::endl;
        return -1;
    }

    // 2. 反序列生成engine
    // 加载了保存在硬盘上的模型文件
    // 存储到std::vector<uchar>类型的engine_data变量中,以便于后续的模型反序列化操作。
    std::vector<uchar> engine_data = {
    
     '\0' };
    load_engine_file(engine_file, engine_data);
    // 使用IRuntime的deserializeCudaEngine方法将其反序列化为ICudaEngine对象。
    // 在TensorRT中,推理的执行是通过ICudaEngine对象来进行的。
    auto mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(engine_data.data(), engine_data.size()));
    if (!mEngine)
    {
    
    
        std::cerr << "Failed to create engine." << std::endl;
        return -1;
    }

    // 3. 创建执行上下文
    // 在 TensorRT 中,推理的执行是通过执行上下文 (ExecutionContext) 来进行的。
    // ExecutionContext 封装了推理运行时所需要的所有信息,包括存储设备上的缓冲区地址、内核参数以及其他设备信息。
    // 因此,当需要在设备上执行模型推理时,ExecutionContext。
    auto context = std::unique_ptr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
    if (!context)
    {
    
    
        std::cerr << "Failed to create ExcutionContext." << std::endl;
        return -1;
    }
    // 4. 创建输入输出缓冲区
    // 在 TensorRT 中,BufferManager 是一个辅助类,用于简化输入/输出缓冲区的创建和管理。
    // TensorRT 提供的 BufferManager 类用于简化这个过程,可以自动创建并管理这些缓冲区,使得在进行推理时不需要手动创建和管理这些缓冲区。
    // 当需要在 GPU 上进行推理时,只需要将输入数据放入 BufferManager 管理的缓冲区中,
    // 然后调用推理函数,等待推理完成后,从 BufferManager 管理的缓冲区中获取推理结果即可。
    samplesCommon::BufferManager buffers(mEngine);


    // 5.读入视频
    auto cap = cv::VideoCapture(input_path_file);
    int width = int(cap.get(cv::CAP_PROP_FRAME_WIDTH));
    int height = int(cap.get(cv::CAP_PROP_FRAME_HEIGHT));
    int fps = int(cap.get(cv::CAP_PROP_FPS));

    // 写入MP4文件,参数分别是:文件名,编码格式,帧率,帧大小
    cv::VideoWriter writer("./output/test.mp4", cv::VideoWriter::fourcc('H', '2', '6', '4'), fps, cv::Size(width, height));

    cv::Mat frame;
    int frame_index = 0;
    // 申请gpu内存
    cudaMalloc(height * width);


    while (cap.isOpened())
    {
    
    
        // 统计运行时间
        auto start = std::chrono::high_resolution_clock::now();

        cap >> frame;
        if (frame.empty())
        {
    
    
            break;    
        }
        frame_index++;
        // 输入预处理(实现了对输入图像处理的gpu 加速)
        process_input(frame, (float*)buffers.getDeviceBuffer(kInputTensorName));
        // 6. 执行推理
        context->executeV2(buffers.getDeviceBindings().data());
        // 推理完成后拷贝回缓冲区
        buffers.copyOutputToHost();

        // 从buffer中提取推理结果
        int32_t* num_det = (int32_t*)buffers.getHostBuffer(kOutNumDet); // 目标检测到的个数
        int32_t* cls = (int32_t*)buffers.getHostBuffer(kOutDetCls);     // 目标检测到的目标类别
        float* conf = (float*)buffers.getHostBuffer(kOutDetScores);     // 目标检测的目标置信度
        float* bbox = (float*)buffers.getHostBuffer(kOutDetBBoxes);     // 目标检测到的目标框
        // 非极大值抑制,得到最后的检测框
        std::vector<Detection> bboxs;
        yolo_nms(bboxs, num_det, cls, conf, bbox, kConfThresh, kNmsThresh);

        // 结束时间
        auto end = std::chrono::high_resolution_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
        auto time_str = std::to_string(elapsed) + "ms";
        auto fps_str = std::to_string(1000 / elapsed) + "fps";

        // 绘制结果
        for (size_t i = 0; i < bboxs.size(); i++)
        {
    
    
            cv::Rect rect = get_rect(frame, bboxs[i].bbox);
            cv::rectangle(frame, rect, cv::Scalar(0x27, 0xC1, 0x36), 2);
            cv::putText(frame, std::to_string((int)bboxs[i].class_id), cv::Point(rect.x, rect.y - 10), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0x27, 0xC1, 0x36), 2);
        }
        cv::putText(frame, time_str, cv::Point(50, 50), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
        cv::putText(frame, fps_str, cv::Point(50, 100), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
        writer.write(frame);
        
        std::cout << "处理完第" << frame_index << "帧" << std::endl;
        if (cv::waitKey(1) == 27)
            break;
    }
    std::cout << "处理完成!!" << std::endl;
    // 程序退出,释放GPU资源,其他均由智能指针自动释放!
    cuda_preprocess_destroy();
    return 0;
}

总结

通过这篇文章,我们开始深入了解了TensorRT,它是一个高性能的深度学习推理优化器和运行时库,主要用于加速深度学习模型的部署。我们讨论了TensorRT的主要优点和工作原理。接下来我们会继续完善TensorRT量化与CUDA编程基础。如果你喜欢我们的文章或者需要源代码全文,可以关注VX公纵号:01编程小屋,发送tensorrt获取源代码全文。

猜你喜欢

转载自blog.csdn.net/weixin_43654363/article/details/131785192