devbox
凡人多烦事01
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

从零到一，激活GPU的力量：使用TensorRT优化和执行深度学习模型，你的TensorRT入门指南_tensorrt improve gpu-util

作者：凡人多烦事01 | 2024-03-18 07:42:51

赞

踩

tensorrt improve gpu-util

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;

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10

IBuilder 是 TensorRT 中的一个接口，负责创建一个优化的 ICudaEngine。ICudaEngine 是一个优化后的可执行网络，可以用于实际执行推理任务。几乎在所有使用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;
}

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10

创建 network 是因为我们需要一个 INetworkDefinition 对象来描述我们的模型。我们使用 builder 的 createNetworkV2 方法来创建一个网络，并为其提供一组标志来指定网络的行为。在这个例子中，我们设置了 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));

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14

在许多神经网络模型中，可能会有多个输入节点。例如，在一个同时接受图像和元数据输入的模型中，可能有一个输入节点用于接收图像，另一个输入节点用于接收元数据。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);

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14

在这个步骤中，我们创建一个 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);

  • 1
  • 2
  • 3

这里我们创建一个 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();

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6

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

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

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

  • 1

这行代码中，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;
}

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
  • 48
  • 49
  • 50
  • 51
  • 52
  • 53
  • 54
  • 55
  • 56
  • 57
  • 58
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68
  • 69
  • 70
  • 71
  • 72
  • 73
  • 74
  • 75
  • 76
  • 77
  • 78
  • 79
  • 80
  • 81
  • 82
  • 83
  • 84
  • 85
  • 86
  • 87
  • 88
  • 89
  • 90
  • 91
  • 92
  • 93
  • 94
  • 95
  • 96
  • 97
  • 98
  • 99
  • 100

TensorRT模型推理

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

  1. 创建推理运行时的runtime：IRuntime 是 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的工作流程中，一般有两个主要阶段：优化阶段和推理阶段。

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 在优化阶段，我们使用IBuilder和INetworkDefinition等接口创建和优化模型。优化完成后，我们通常将优化后的模型序列化并保存到硬盘上。
  • 在推理阶段，我们需要从硬盘上加载优化后的模型，然后执行推理。这个阶段就需要用到IRuntime。我们首先使用IRuntime的deserializeCudaEngine方法从序列化的数据中加载模型，然后使用加载的模型进行推理。

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

  1. 读取模型文件并反序列生成engine与创建ICudaEngine神经网络模型对象：我们首先加载了保存在硬盘上的模型文件，然后使用IRuntime的deserializeCudaEngine方法将其反序列化为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();// 关闭文件
}

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20

反序列读取后实例化为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;
}

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11

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

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

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

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

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

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7

在 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);

  • 1
  • 2

在上述代码中，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
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13

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

  1. 打开一个视频文件并获取其属性：cv::VideoCapture对象被用来打开输入的视频文件。通过调用cv::CAP_PROP_FRAME_WIDTH，cv::CAP_PROP_FRAME_HEIGHT和cv::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
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14

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

  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函数进行非极大值抑制，这是目标检测后处理的常见步骤，用于消除冗余的检测框，只保留最佳的检测结果。kConfThresh和kNmsThresh是非极大值抑制的两个阈值，一般会通过实验来确定。

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

在如上代码中，会有一些特别的变量，通常由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
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47

上面的变量中，通常用来表示：

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

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

主要包含了一个描述数据框的数据结构Detection和两个函数iou和yolo_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;
                }
            }
        }
    }
}

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
  • 48
  • 49
  • 50
  • 51
  • 52
  • 53
  • 54
  • 55
  • 56
  • 57
  • 58
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68
  • 69
  • 70
  • 71
  • 72
  • 73
  • 74
  • 75
  • 76
  • 77
  • 78
  • 79
  • 80
  • 81
  • 82
  • 83
  • 84
  • 85
  • 86

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;
}

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
  • 48
  • 49
  • 50
  • 51
  • 52
  • 53
  • 54
  • 55
  • 56
  • 57
  • 58
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68
  • 69
  • 70
  • 71
  • 72
  • 73
  • 74
  • 75
  • 76
  • 77
  • 78
  • 79
  • 80
  • 81
  • 82
  • 83
  • 84
  • 85
  • 86
  • 87
  • 88
  • 89
  • 90
  • 91
  • 92
  • 93
  • 94
  • 95
  • 96
  • 97
  • 98
  • 99
  • 100
  • 101
  • 102
  • 103
  • 104
  • 105
  • 106
  • 107
  • 108
  • 109
  • 110
  • 111
  • 112
  • 113
  • 114
  • 115
  • 116
  • 117
  • 118
  • 119
  • 120
  • 121
  • 122
  • 123
  • 124
  • 125
  • 126
  • 127
  • 128
  • 129
  • 130
  • 131
  • 132
  • 133
  • 134
  • 135
  • 136
  • 137
  • 138
  • 139
  • 140
  • 141
  • 142
  • 143
  • 144
  • 145
  • 146
  • 147
  • 148
  • 149
  • 150
  • 151
  • 152
  • 153
  • 154
  • 155
  • 156
  • 157
  • 158
  • 159
  • 160
  • 161
  • 162
  • 163
  • 164
  • 165
  • 166
  • 167
  • 168
  • 169
  • 170
  • 171
  • 172
  • 173
  • 174
  • 175
  • 176
  • 177
  • 178
  • 179
  • 180
  • 181
  • 182
  • 183
  • 184
  • 185
  • 186
  • 187
  • 188
  • 189
  • 190
  • 191
  • 192
  • 193
  • 194
  • 195
  • 196
  • 197
  • 198
  • 199
  • 200
  • 201
  • 202
  • 203
  • 204
  • 205
  • 206
  • 207
  • 208
  • 209
  • 210
  • 211
  • 212
  • 213
  • 214
  • 215
  • 216
  • 217
  • 218
  • 219
  • 220
  • 221
  • 222
  • 223
  • 224
  • 225
  • 226
  • 227
  • 228
  • 229
  • 230
  • 231
  • 232
  • 233
  • 234
  • 235
  • 236
  • 237
  • 238
  • 239
  • 240
  • 241
  • 242
  • 243
  • 244
  • 245
  • 246
  • 247
  • 248
  • 249
  • 250
  • 251
  • 252
  • 253
  • 254
  • 255
  • 256
  • 257
  • 258
  • 259
  • 260
  • 261

总结

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

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/凡人多烦事01/article/detail/261425
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号