De zero a um, ative o poder da GPU: use o TensorRT para otimizar e executar modelos de aprendizado profundo, seu guia de primeiros passos do TensorRT

Notas de estudo do TensorRT

Otimização do modelo TensorRT

Otimização de modelo e introdução de processo

O TensorRT é um otimizador de inferência de aprendizagem profunda (Inferência) de alto desempenho e uma biblioteca de tempo de execução fornecida pela NVIDIA que pode ser usada para acelerar modelos de aprendizagem profunda. O principal objetivo do TensorRT é otimizar o uso da computação paralela para utilizar totalmente o poder de computação das GPUs NVIDIA. Ele funciona com vários tipos de modelos de aprendizado profundo, incluindo CNN (Convolutional Neural Network), RNN (Recurrent Neural Network), etc.

Os métodos e estratégias para o TensorRT implementar a otimização do modelo incluem:

  1. Fusão de camada e tensor : ao fundir camadas sucessivas de uma rede, o TensorRT pode reduzir o número de leituras e gravações de dados da memória da GPU, melhorando assim o desempenho.
  2. Ajuste de precisão : o TensorRT suporta cálculos de precisão mista e você pode escolher FP32, FP16 ou INT8 para cálculos para equilibrar precisão e desempenho.
  3. Tamanho de entrada dinâmico : o TensorRT suporta tamanho de entrada dinâmico, o que significa que você pode executar inferência no mesmo modelo com entradas de tamanhos diferentes.
  4. Agendamento automático de kernel : o TensorRT selecionará automaticamente o kernel CUDA ideal para cálculo, o que permite o melhor desempenho em diferentes hardwares.

O processo geral geralmente inclui as seguintes etapas:

  1. Criar e configurar construtor : use para nvinfer1::createInferBuilder()criar um objeto construtor para construção de modelo subseqüente.
  2. Criar e configurar rede : use para nvinfer1::createNetworkV2()criar um objeto de rede para descrever o modelo.
  3. Definir tamanho de entrada do modelo : defina o intervalo possível de entrada do modelo de acordo com suas necessidades, o que ajuda o TensorRT a otimizar o modelo.
  4. Criar e configurar configuração : use para nvinfer1::createBuilderConfig()criar um objeto de configuração para salvar várias configurações durante a otimização do modelo, como precisão, tamanho máximo do lote, etc.
  5. Crie um fluxo CUDA e defina o fluxo de perfil de configuração : Use para makeCudaStream()criar um fluxo CUDA e, em seguida, use setProfileStream()para configurá-lo para configurar operações de GPU assíncronas.
  6. Construir e serializar o modelo : use para buildSerializedNetwork()criar e otimizar o modelo e, em seguida, serializá-lo para um processo de inferência posterior.

Codificação da etapa do processo de otimização do modelo

  1. Crie e configure o construtor: primeiro crie um IBuilder, então use isso IBuilderpara criar um 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é uma interface no TensorRT responsável por criar um arquivo ICudaEngine. ICudaEngineé uma rede executável otimizada que pode ser usada para executar tarefas de inferência. Ele será usado em quase todas as ocasiões em que o TensorRT for usado , pois é uma ferramenta fundamental para IBuildera criação de mecanismos de execução otimizados ( ). ICudaEngineNão importa qual seja o seu modelo, desde que você queira usar o TensorRT para otimização e implantação, você precisa criá-lo e usá-lo IBuilder. Portanto, criar IBuilderé uma etapa necessária, pois é ICudaEngineuma ferramenta fundamental para criar arquivos . Aqui é usado std::unique_ptrpara gerenciar IBuildera memória do arquivo .

Um objeto de log precisa ser fornecido no momento da criação para receber erros, avisos e outras mensagens. Esta etapa é o ponto de partida de todo o processo de criação do modelo, sem a qual buildernão podemos criar a rede (modelo).

  1. Criar e configurar a rede
// ========== 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;
}

Criado networkporque precisamos de um INetworkDefinitionobjeto para descrever nosso modelo. Usamos o método builderde createNetworkV2para criar uma rede e fornecer a ela um conjunto de sinalizadores para especificar o comportamento da rede. Neste exemplo, definimos kEXPLICIT_BATCHo sinalizador para habilitar lotes explícitos.

BatchNormalmente usado quando você precisa fazer inferência em uma grande quantidade de dados. Durante os estágios de treinamento e inferência, geralmente não processamos uma amostra por vez, mas um lote de amostras por vez, o que pode tornar o uso mais eficiente dos recursos de computação, especialmente quando se computa em GPUs. Ao mesmo tempo, escolher um Batchtamanho adequado também é uma questão de equilíbrio. Por um lado, um Batchtamanho maior pode aproveitar melhor o paralelismo de hardware e melhorar a eficiência da computação. Por outro lado, um Batchtamanho maior consumirá mais recursos de memória, então precisa escolher de acordo com a situação específica.

  1. Definir tamanho de entrada do modelo

No TensorRT, o perfil de otimização é usado para descrever o tamanho de entrada e a faixa de tamanho dinâmico do modelo. Ao otimizar o arquivo de configuração, você pode informar ao TensorRT o intervalo de tamanhos possíveis dos dados de entrada, para que ele possa criar um modelo otimizado para vários tamanhos de entrada.

// 配置网络参数
// 我们需要告诉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));

Em muitos modelos de redes neurais, pode haver vários nós de entrada. Por exemplo, em um modelo que aceita entradas de imagem e metadados, pode haver um nó de entrada para imagens e outro nó de entrada para metadados. auto input = network->getInput(0);Esta linha de código obtém o primeiro nó de entrada da rede. O nó de entrada da rede é a camada de entrada do modelo, que recebe os dados de entrada do modelo. Aqui 0está um índice referente ao primeiro nó de entrada da rede. getInputO parâmetro da função é usado para especificar qual nó de entrada você deseja obter. Se você passar 0, o primeiro nó de entrada será retornado. Se você passar 1, o segundo nó de entrada será retornado e assim por diante. Portanto, você pode passar diferentes parâmetros para obter diferentes nós de entrada de acordo com sua estrutura de rede e necessidades. No código que você forneceu, como a rede pode ter apenas um nó de entrada, passe 0para obter o nó de entrada.

Nesta etapa, primeiro obtemos os nós de entrada do modelo (neste exemplo, assumimos que o modelo possui apenas uma entrada) e, em seguida, criamos um perfil de otimização ( OptimizationProfile). Um perfil de otimização descreve o possível intervalo de entradas do modelo, o que é necessário para o processo de otimização do modelo. Definimos o tamanho mínimo, ideal e máximo para cada dimensão.

Nesse código, cada camada de rede tem um nome exclusivo e getNameo método é usado para obter o nome do nó de entrada (ou mais geralmente, o tensor). Esse nome é fornecido ao definir o modelo de rede e geralmente é usado para ajudar a identificar e rastrear nós individuais na rede. Esse nome tem um uso importante no TensorRT, pois é usado ao definir o tamanho dos nós de entrada e saída ou ao realizar inferência.

  1. Criar e configurar configuração
// ========== 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);

Nesta etapa, criamos um IBuilderConfigobjeto para salvar várias configurações durante a otimização do modelo. Adicionamos a ele o perfil de otimização que criamos anteriormente e definimos alguns sinalizadores como , kFP16para habilitar o suporte para precisão FP16. Também definimos um tamanho máximo de lote como parte do processo de otimização.

No TensorRT, os arquivos de configuração ( IBuilderConfigobjetos) são usados ​​para salvar várias configurações durante a otimização do modelo. Essas configurações incluem:

  • Modo de precisão otimizado: você pode escolher o modo de precisão FP32, FP16 ou INT8. O uso de menor precisão (como FP16 ou INT8) pode reduzir o uso de recursos de computação e, assim, melhorar a velocidade de inferência, mas pode sacrificar alguma precisão de inferência.
  • Tamanho máximo do lote: este é o número máximo de lotes de entrada que podem ser processados ​​durante a otimização do modelo.
  • Tamanho do espaço de trabalho: esta é a quantidade máxima de memória que o TensorRT pode usar ao otimizar e executar o modelo.
  • Perfil de otimização: um perfil de otimização descreve o possível intervalo de dimensões de entrada do modelo. A partir dessas informações, o TensorRT pode criar um modelo otimizado para vários tamanhos de entrada.

Os arquivos de configuração também podem conter outras configurações, como seleção de dispositivo de GPU, seleção de política de camada e assim por diante. Essas configurações afetarão como o TensorRT otimiza o modelo e o desempenho do modelo otimizado.

  1. Crie um fluxo CUDA e defina o fluxo de perfil de configuração
// 创建流,用于设置profile
auto profileStream = samplesCommon::makeCudaStream();
config->setProfileStream(*profileStream);

Aqui criamos um fluxo CUDA para executar operações de GPU de forma assíncrona. Em seguida, definimos esse fluxo CUDA como o fluxo de perfil do objeto de configuração. O fluxo de perfil é usado para determinar quando o tempo do kernel pode ser coletado e quando o tempo coletado pode ser lido.

Aqui profileStreamestá o fluxo CUDA (CUDA Stream), não o perfil de otimização anterior (perfil de otimização). Um fluxo CUDA é uma fila ordenada de operações que são executadas de forma assíncrona na GPU. Os fluxos podem ser usados ​​para organizar e controlar a simultaneidade e as dependências de execução.

No seu código, profileStreamé samplesCommon::makeCudaStream()criado via . Este fluxo será passado para o objeto de configuração e, em seguida, o TensorRT executará as operações relacionadas à configotimização do arquivo de configuração () neste fluxo . profileIsso permite que essas operações sejam executadas simultaneamente, melhorando a eficiência da execução.

Deve-se observar que o fluxo CUDA aqui e o arquivo de configuração de otimização anterior são dois conceitos completamente diferentes. O arquivo de configuração de otimização contém as informações do tamanho de entrada do modelo, que é usado para orientar o processo de otimização do modelo; enquanto o fluxo CUDA é usado para gerenciar a ordem de execução das operações da GPU para melhorar a eficiência da execução. Embora os dois tenham nomes semelhantes, eles são completamente diferentes em função e propósito.

  1. Construir e serializar o modelo
// ========== 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();

Nesta etapa final, usamos builder, , networke configpara construir e serializar nosso modelo. Este processo pode realizar uma série de otimizações no modelo, incluindo fusão de camadas, seleção de algoritmo de convolução, seleção de formato de tensor, etc. Esta etapa produz um modelo serializado que podemos salvar em um arquivo e, em seguida, carregar e usar posteriormente no processo de inferência. Isso pode reduzir bastante o tempo de carregamento do modelo, pois não precisamos passar pelo processo de otimização novamente.

Construa e serialize um modelo otimizado usando os objetos builder, networke previamente criados e configurados. A função retorna um objeto que contém os dados do modelo serializado.configbuildSerializedNetworkIHostMemory

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

Nesta linha de código planhá um ponteiro exclusivo para a memória dos dados do modelo serializado.

O principal objetivo deste código é criar um modelo serializado e salvá-lo em um arquivo. A vantagem de serializar e salvar o modelo em um arquivo é que você pode carregar esse arquivo diretamente em execuções subsequentes do programa sem ter que reconstruir e otimizar o modelo todas as vezes, o que pode melhorar muito a eficiência.

Notas de ponto de conhecimento

  • IBuilderé INetworkDefinitiono objeto usado para criar o , que é o objeto usado para criar o modelo.
  • INetworkDefinitionRepresenta uma rede TensorRT usada para descrever a estrutura do modelo.
  • Lote explícito significa que o tamanho do lote é explicitamente especificado quando o modelo é construído, em vez de determinado em tempo de execução. Isso pode melhorar o desempenho do modelo em alguns casos.
  • IBuilderConfigEle é usado para salvar várias configurações durante a otimização do modelo, incluindo modo de precisão de otimização, tamanho máximo do lote, tamanho do espaço de trabalho e arquivo de configuração de otimização, etc.
  • Um perfil de otimização descreve o possível intervalo de dimensões de entrada do modelo e pode orientar o processo de otimização do modelo.
  • Um fluxo CUDA é uma fila ordenada de operações que são executadas de forma assíncrona na GPU. Os fluxos podem ser usados ​​para organizar e controlar a simultaneidade e as dependências de execução.
  • Use buildSerializedNetworko método para criar e serializar um modelo e retornar um IHostMemoryobjeto que contém os dados do modelo serializado. Serializar e salvar o modelo em um arquivo pode melhorar a velocidade de carregamento do modelo.
  • No TensorRT, todas as entradas do modelo são camadas de rede, e cada camada de rede possui um nome exclusivo, que pode ser getNameobtido usando o método.
TensorRT模型优化
用于创建 INetworkDefinition 对象的接口
用于描述模型的对象
设置模型输入的可能范围,用于模型优化
保存模型优化过程中的各种设置
用于异步执行 GPU 操作,确定何时可以开始收集内核时间,何时可以读取收集到的时间
进行模型优化,生成序列化模型,用于之后的推理过程
开始
创建并配置 builder
创建并配置 network
设置模型输入尺寸
创建并配置 config
创建 CUDA 流并设置 config 的 profile stream
构建并序列化模型
结束

código completo

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

Inferência de modelo TensorRT

Implantação do modelo e codificação da etapa do processo de inferência

  1. Runtime para criação de inferência runtime : IRuntimeÉ uma interface fornecida pelo TensorRT, que é usada principalmente para executar modelos serializados na fase de inferência. Criar uma instância é a primeira etapaIRuntime para carregar e executar o mecanismo TensorRT durante a inferência . Se você precisar criar uma instância. Function é uma função fornecida pelo TensorRT para criar instâncias. Esta função requer uma instância como parâmetro, que é usada principalmente para processar as informações de log do TensorRT durante a operação.IRuntimecreateInferRuntimeIRuntimeILoggerILogger
// 使用 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的工作流程中,一般有两个主要阶段:优化阶段和推理阶段。
  • Na fase de otimização, criamos e otimizamos o modelo usando interfaces como IBuildere INetworkDefinition. Após a conclusão da otimização, geralmente serializamos e salvamos o modelo otimizado no disco.
  • No estágio de inferência, precisamos carregar o modelo otimizado do disco rígido e então realizar a inferência. É necessário nesta fase IRuntime. IRuntimeO método que usamos primeiro deserializeCudaEnginecarrega o modelo dos dados serializados e, em seguida, usa o modelo carregado para inferência.

所以,只要需要使用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获取源代码全文。

Acho que você gosta

Origin blog.csdn.net/weixin_43654363/article/details/131785192
Recomendado
Clasificación