Aktivieren Sie von Null auf Eins die Leistung der GPU: Verwenden Sie TensorRT, um Deep-Learning-Modelle zu optimieren und auszuführen, Ihren TensorRT-Einstiegsleitfaden

Anmerkungen zur TensorRT-Studie

TensorRT-Modelloptimierung

Modelloptimierung und Prozesseinführung

TensorRT ist ein leistungsstarker Deep-Learning-Inferenz-(Inferenz-)Optimierer und eine Laufzeitbibliothek von NVIDIA, die zur Beschleunigung von Deep-Learning-Modellen verwendet werden kann. Das Hauptziel von TensorRT besteht darin, den Einsatz paralleler Datenverarbeitung zu optimieren, um die Rechenleistung von NVIDIA-GPUs voll auszunutzen. Es funktioniert mit verschiedenen Arten von Deep-Learning-Modellen, einschließlich CNN (Convolutional Neural Network), RNN (Recurrent Neural Network) usw.

Zu den Methoden und Strategien für TensorRT zur Implementierung der Modelloptimierung gehören:

  1. Layer- und Tensor-Fusion : Durch die Fusion aufeinanderfolgender Schichten eines Netzwerks kann TensorRT die Anzahl der Datenlese- und -schreibvorgänge aus dem GPU-Speicher reduzieren und dadurch die Leistung verbessern.
  2. Präzisionsanpassung : TensorRT unterstützt Berechnungen mit gemischter Genauigkeit, und Sie können FP32, FP16 oder INT8 für Berechnungen wählen, um Präzision und Leistung in Einklang zu bringen.
  3. Dynamische Eingabegröße : TensorRT unterstützt die dynamische Eingabegröße, was bedeutet, dass Sie Rückschlüsse auf demselben Modell mit Eingaben unterschiedlicher Größe ausführen können.
  4. Automatische Kernel-Planung : TensorRT wählt automatisch den optimalen CUDA-Kernel für die Berechnung aus, was die beste Leistung auf unterschiedlicher Hardware ermöglicht.

Der allgemeine Prozess umfasst im Allgemeinen die folgenden Schritte:

  1. Builder erstellen und konfigurieren : Zum nvinfer1::createInferBuilder()Erstellen eines Builder-Objekts für die anschließende Modellerstellung.
  2. Netzwerk erstellen und konfigurieren : Zum nvinfer1::createNetworkV2()Erstellen eines Netzwerkobjekts zur Beschreibung des Modells.
  3. Modelleingabegröße festlegen : Stellen Sie den möglichen Bereich der Modelleingabe entsprechend Ihren Anforderungen ein, was TensorRT bei der Optimierung des Modells hilft.
  4. Konfiguration erstellen und konfigurieren : Zum nvinfer1::createBuilderConfig()Erstellen eines Konfigurationsobjekts zum Speichern verschiedener Einstellungen während der Modelloptimierung, z. B. Genauigkeit, maximale Batchgröße usw.
  5. Erstellen Sie einen CUDA-Stream und legen Sie den Profil-Stream von „config“ fest : Erstellen Sie makeCudaStream()einen CUDA-Stream und setProfileStream()legen Sie ihn dann auf „config“ für asynchrone GPU-Vorgänge fest.
  6. Erstellen und serialisieren Sie das Modell : Verwenden Sie es, um buildSerializedNetwork()das Modell zu erstellen und zu optimieren und es dann für einen späteren Inferenzprozess zu serialisieren.

Codierung der Schritte des Modelloptimierungsprozesses

  1. Erstellen und konfigurieren Sie den Builder: Erstellen Sie zuerst einen IBuilderund verwenden Sie dann diesen, IBuilderum einen zu erstellen 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;

IBuilderist eine Schnittstelle in TensorRT, die für die Erstellung einer optimierten ICudaEngine. ICudaEngineist ein optimiertes ausführbares Netzwerk, das zur tatsächlichen Ausführung von Inferenzaufgaben verwendet werden kann. Es wird fast überall dort zum Einsatz kommen, wo TensorRT zum Einsatz kommt , da es ein Schlüsselwerkzeug für IBuilderdie Erstellung optimierter Ausführungs-Engines ist ( ). ICudaEngineUnabhängig von Ihrem Modell müssen Sie TensorRT erstellen und verwenden, solange Sie es zur Optimierung und Bereitstellung verwenden möchten IBuilder. Daher IBuilderist das Erstellen ein notwendiger Schritt, da es ICudaEngineein wichtiges Werkzeug zum Erstellen darstellt. Hier wird der Speicher des . std::unique_ptrverwaltet .IBuilder

Zum Zeitpunkt der Erstellung muss ein Protokollobjekt bereitgestellt werden, um Fehler, Warnungen und andere Meldungen zu empfangen. Dieser Schritt ist der Ausgangspunkt des gesamten Modellerstellungsprozesses, ohne den builderwir das Netzwerk (Modell) nicht erstellen können.

  1. Erstellen und konfigurieren Sie das Netzwerk
// ========== 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;
}

Erstellt network, weil wir ein INetworkDefinitionObjekt benötigen, um unser Modell zu beschreiben. Wir verwenden die Methode builder, um createNetworkV2ein Netzwerk zu erstellen und es mit einer Reihe von Flags zu versehen, die das Verhalten des Netzwerks festlegen. In diesem Beispiel setzen wir kEXPLICIT_BATCHdas Flag, um die explizite Stapelverarbeitung zu ermöglichen.

BatchWird normalerweise verwendet, wenn Sie Rückschlüsse auf eine große Datenmenge ziehen müssen. Während der Trainings- und Inferenzphase verarbeiten wir normalerweise nicht eine Probe nach der anderen, sondern eine Reihe von Proben nach der anderen, wodurch die Rechenressourcen effizienter genutzt werden können, insbesondere bei der Berechnung auf GPUs. Gleichzeitig Batchist die Wahl einer geeigneten Größe auch eine Frage des Gleichgewichts. Einerseits Batchkann eine größere Größe die Hardware-Parallelität besser nutzen und die Recheneffizienz verbessern. Andererseits Batchverbraucht eine größere Größe mehr Speicherressourcen. Daher muss entsprechend der jeweiligen Situation eine Auswahl getroffen werden.

  1. Legen Sie die Modelleingabegröße fest

In TensorRT wird das Optimierungsprofil verwendet, um die Eingabegröße und den dynamischen Größenbereich des Modells zu beschreiben. Durch die Optimierung der Konfigurationsdatei können Sie TensorRT den Bereich möglicher Größen der Eingabedaten mitteilen, sodass ein optimiertes Modell für verschiedene Eingabegrößen erstellt werden kann.

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

In vielen neuronalen Netzwerkmodellen kann es mehrere Eingabeknoten geben. Beispielsweise kann es in einem Modell, das sowohl Bild- als auch Metadateneingaben akzeptiert, einen Eingabeknoten für Bilder und einen anderen Eingabeknoten für Metadaten geben. auto input = network->getInput(0);Diese Codezeile ruft den ersten Eingabeknoten des Netzwerks ab. Der Eingabeknoten des Netzwerks ist die Eingabeschicht des Modells, die die Eingabedaten des Modells empfängt. Hier 0ist ein Index, der sich auf den ersten Eingabeknoten des Netzwerks bezieht. getInputDer Parameter der Funktion wird verwendet, um anzugeben, welchen Eingabeknoten Sie erhalten möchten. Wenn Sie übergeben 0, wird der erste Eingabeknoten zurückgegeben. Wenn Sie übergeben 1, wird der zweite Eingabeknoten zurückgegeben und so weiter. Daher können Sie verschiedene Parameter übergeben, um entsprechend Ihrer Netzwerkstruktur und Ihren Anforderungen unterschiedliche Eingabeknoten zu erhalten. Da das Netzwerk möglicherweise nur einen Eingabeknoten hat, übergeben Sie in dem von Ihnen angegebenen Code 0den Eingabeknoten.

In diesem Schritt rufen wir zunächst die Eingabeknoten des Modells ab (in diesem Beispiel gehen wir davon aus, dass das Modell nur eine Eingabe hat) und erstellen dann ein Optimierungsprofil ( OptimizationProfile). Ein Optimierungsprofil beschreibt den möglichen Bereich an Modelleingaben, der für den Optimierungsprozess des Modells notwendig ist. Wir legen für jede Dimension die minimale, optimale und maximale Größe fest.

In diesem Code hat jede Netzwerkschicht einen eindeutigen Namen und getNamedie Methode wird verwendet, um den Namen des Eingabeknotens (oder allgemeiner des Tensors) abzurufen. Dieser Name wird bei der Definition des Netzwerkmodells angegeben und dient normalerweise dazu, einzelne Knoten im Netzwerk zu identifizieren und zu verfolgen. Dieser Name hat in TensorRT eine wichtige Verwendung, da er beim Festlegen der Größe der Eingabe- und Ausgabeknoten oder beim Durchführen von Inferenzen verwendet wird.

  1. Konfiguration erstellen und konfigurieren
// ========== 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);

In diesem Schritt erstellen wir ein IBuilderConfigObjekt, um verschiedene Einstellungen während der Modelloptimierung zu speichern. Wir fügen das zuvor erstellte Optimierungsprofil hinzu und setzen einige Flags wie , kFP16um die Unterstützung für FP16-Präzision zu ermöglichen. Im Rahmen der Optimierung legen wir auch eine maximale Losgröße fest.

In TensorRT werden Konfigurationsdateien ( IBuilderConfigObjekte) verwendet, um verschiedene Einstellungen während der Modelloptimierung zu speichern. Zu diesen Einstellungen gehören:

  • Optimierter Präzisionsmodus: Sie können den Präzisionsmodus FP32, FP16 oder INT8 wählen. Die Verwendung einer geringeren Präzision (z. B. FP16 oder INT8) kann den Einsatz von Rechenressourcen reduzieren und somit die Inferenzgeschwindigkeit verbessern, kann jedoch zu Einbußen bei der Inferenzgenauigkeit führen.
  • Maximale Batchgröße: Dies ist die maximale Anzahl von Eingabebatches, die während der Modelloptimierung verarbeitet werden können.
  • Arbeitsbereichsgröße: Dies ist die maximale Speichermenge, die TensorRT beim Optimieren und Ausführen des Modells verwenden kann.
  • Optimierungsprofil: Ein Optimierungsprofil beschreibt den möglichen Bereich der Modelleingabedimensionen. Aus diesen Informationen kann TensorRT ein für verschiedene Eingabegrößen optimiertes Modell erstellen.

Konfigurationsdateien können auch andere Einstellungen wie die Auswahl des GPU-Geräts, die Auswahl der Layer-Richtlinie usw. enthalten. Diese Einstellungen wirken sich darauf aus, wie TensorRT das Modell optimiert und wie die Leistung des optimierten Modells ist.

  1. Erstellen Sie einen CUDA-Stream und legen Sie den Profil-Stream der Konfiguration fest
// 创建流,用于设置profile
auto profileStream = samplesCommon::makeCudaStream();
config->setProfileStream(*profileStream);

Hier erstellen wir einen CUDA-Stream, um GPU-Operationen asynchron auszuführen. Dann legen wir diesen CUDA-Stream als Profil-Stream des Konfigurationsobjekts fest. Der Profilstream wird verwendet, um zu bestimmen, wann Kernelzeit erfasst werden kann und wann die erfasste Zeit gelesen werden kann.

Hier profileStreamist der CUDA-Stream (CUDA Stream), nicht das vorherige Optimierungsprofil (Optimierungsprofil). Ein CUDA-Stream ist eine geordnete Warteschlange von Vorgängen, die asynchron auf der GPU ausgeführt werden. Streams können verwendet werden, um die Parallelität und Abhängigkeiten der Ausführung zu organisieren und zu steuern.

In Ihrem Code profileStreamwird samplesCommon::makeCudaStream()über erstellt. Dieser Stream wird an das Konfigurationsobjekt übergeben , und dann führt TensorRT Vorgänge im Zusammenhang mit configder Optimierung der Konfigurationsdatei () in diesem Stream aus. profileDadurch können diese Vorgänge gleichzeitig ausgeführt werden, was die Ausführungseffizienz verbessert.

Es ist zu beachten, dass der CUDA-Stream hier und die vorherige Optimierungskonfigurationsdatei zwei völlig unterschiedliche Konzepte sind. Die Optimierungskonfigurationsdatei enthält Informationen zur Modelleingabegröße, die zur Steuerung des Optimierungsprozesses des Modells verwendet werden, während der CUDA-Stream zur Verwaltung der Ausführungsreihenfolge von GPU-Operationen verwendet wird, um die Ausführungseffizienz zu verbessern. Obwohl die beiden ähnliche Namen haben, unterscheiden sie sich tatsächlich in Funktion und Zweck völlig.

  1. Erstellen und serialisieren Sie das Modell
// ========== 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();

In diesem letzten Schritt verwenden wir builder, networkund , configum unser Modell zu erstellen und zu serialisieren. Dieser Prozess kann eine Reihe von Optimierungen am Modell durchführen, einschließlich Schichtfusion, Auswahl des Faltungsalgorithmus, Auswahl des Tensorformats usw. Dieser Schritt erzeugt ein serialisiertes Modell, das wir in einer Datei speichern und dann laden und später im Inferenzprozess verwenden können. Dies kann die Ladezeit des Modells erheblich verkürzen, da wir den Optimierungsprozess nicht erneut durchlaufen müssen.

builderErstellen und serialisieren Sie ein optimiertes Modell unter Verwendung der zuvor erstellten und konfigurierten networkObjekte . Die Funktion gibt ein Objekt zurück, das die serialisierten Modelldaten enthält.configbuildSerializedNetworkIHostMemory

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

In dieser Codezeile planbefindet sich ein eindeutiger Zeiger auf den Speicher der serialisierten Modelldaten.

Der Hauptzweck dieses Codes besteht darin, ein serialisiertes Modell zu erstellen und es in einer Datei zu speichern. Der Vorteil der Serialisierung und Speicherung des Modells in einer Datei besteht darin, dass Sie diese Datei in nachfolgenden Programmläufen direkt laden können, ohne das Modell jedes Mal neu erstellen und optimieren zu müssen, was die Effizienz erheblich verbessern kann.

Notizen zu Wissenspunkten

  • IBuilderist INetworkDefinitiondas zum Erstellen des Modells verwendete Objekt.
  • INetworkDefinitionStellt ein TensorRT-Netzwerk dar, das zur Beschreibung der Struktur des Modells verwendet wird.
  • Explizite Batchverarbeitung bedeutet, dass die Batchgröße beim Erstellen des Modells explizit angegeben und nicht zur Laufzeit bestimmt wird. Dies kann in einigen Fällen die Leistung des Modells verbessern.
  • IBuilderConfigEs wird verwendet, um verschiedene Einstellungen während der Modelloptimierung zu speichern, einschließlich Optimierungsgenauigkeitsmodus, maximale Stapelgröße, Arbeitsbereichsgröße und Optimierungskonfigurationsdatei usw.
  • Ein Optimierungsprofil beschreibt den möglichen Bereich der Modelleingabedimensionen und kann den Optimierungsprozess des Modells steuern.
  • Ein CUDA-Stream ist eine geordnete Warteschlange von Vorgängen, die asynchron auf der GPU ausgeführt werden. Streams können verwendet werden, um die Parallelität und Abhängigkeiten der Ausführung zu organisieren und zu steuern.
  • Verwenden Sie buildSerializedNetworkdie Methode, um ein Modell zu erstellen und zu serialisieren und ein IHostMemoryObjekt zurückzugeben, das die serialisierten Modelldaten enthält. Durch das Serialisieren und Speichern des Modells in einer Datei kann die Ladegeschwindigkeit des Modells verbessert werden.
  • In TensorRT sind alle Modelleingaben Netzwerkschichten, und jede Netzwerkschicht hat einen eindeutigen Namen, der getNamemit der Methode abgerufen werden kann.
TensorRT模型优化
用于创建 INetworkDefinition 对象的接口
用于描述模型的对象
设置模型输入的可能范围,用于模型优化
保存模型优化过程中的各种设置
用于异步执行 GPU 操作,确定何时可以开始收集内核时间,何时可以读取收集到的时间
进行模型优化,生成序列化模型,用于之后的推理过程
开始
创建并配置 builder
创建并配置 network
设置模型输入尺寸
创建并配置 config
创建 CUDA 流并设置 config 的 profile stream
构建并序列化模型
结束

Vollständiger Code

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-Modellinferenz

Schrittkodierung für Modellbereitstellung und Inferenzprozess

  1. Laufzeit zum Erstellen von Inferenzlaufzeiten : IRuntimeHierbei handelt es sich um eine von TensorRT bereitgestellte Schnittstelle, die hauptsächlich zum Ausführen serialisierter Modelle in der Inferenzphase verwendet wird. Das Erstellen einer Instanz ist der erste SchrittIRuntime beim Laden und Ausführen der TensorRT-Engine während der Inferenz . Wenn Sie eine Instanz erstellen müssen . Funktion ist eine von TensorRT bereitgestellte Funktion zum Erstellen von Instanzen. Diese Funktion erfordert eine Instanz als Parameter, die hauptsächlich zur Verarbeitung der Protokollinformationen von TensorRT während des Betriebs verwendet wird.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的工作流程中,一般有两个主要阶段:优化阶段和推理阶段。
  • In der Optimierungsphase erstellen und optimieren wir das Modell mithilfe von Schnittstellen wie IBuilderund INetworkDefinition. Nach Abschluss der Optimierung serialisieren wir normalerweise das optimierte Modell und speichern es auf der Festplatte.
  • In der Inferenzphase müssen wir das optimierte Modell von der Festplatte laden und dann die Inferenz durchführen. Dies ist in dieser Phase erforderlich IRuntime. IRuntimeDie von uns verwendete Methode deserializeCudaEnginelädt zunächst das Modell aus den serialisierten Daten und verwendet dann das geladene Modell zur Inferenz.

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

Ich denke du magst

Origin blog.csdn.net/weixin_43654363/article/details/131785192
Empfohlen
Rangfolge