devbox
我家小花儿
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

TensorRT创建Engine并推理engine_tensorrt构建engine流程图

作者：我家小花儿 | 2024-03-18 07:43:18

赞

踩

tensorrt构建engine流程图

1. 验证集数据集

               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|██████████| 84/
                 all       1000      28423      0.451      0.374      0.376      0.209
         pedestrians       1000      17833      0.737      0.855       0.88      0.609
              riders       1000        185      0.545      0.492      0.521      0.256
partially-visible-person       1000       9335      0.456      0.338      0.336      0.125
      ignore-regions       1000        409       0.37      0.138      0.121     0.0485
               crowd       1000        661      0.146     0.0454     0.0237    0.00837
Speed: 0.1ms pre-process, 1.4ms inference, 1.6ms NMS per image at shape (12, 3, 640, 640)

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

2. 视频流推理

Speed: 0.3ms pre-process, 9.0ms inference, 0.9ms NMS per image at shape (1, 3, 640, 640)

  • 1

计算FPS

要计算FPS（每秒帧数），我们需要将每张图像的时间（包括预处理、推理和NMS）倒数，得到每秒可以处理的图像数量，然后乘以批量大小。

默认batch_size = 1，则每张图像的总时间为0.3毫秒+9.0毫秒+0.9毫秒=10.2毫秒。

因此，FPS将是1 /（10.2毫秒）= 98.04 FPS。

3. 使用TensorRT在Tesla A40 12Q上的推理是

在这里插入图片描述

3. 理解runtime和engine

runtime是TensorRT运行时环境，它提供了创建和管理TensorRT引擎的API。在运行时环境中，可以加载和运行TensorRT引擎，使用TensorRT API管理内存分配和复制，以及在GPU上执行各种TensorRT操作。

engine是一个已经被序列化和优化的TensorRT模型。TensorRT引擎是用NVIDIA TensorRT库创建的，它将深度学习模型转换为可在GPU上高效执行的格式。要使用TensorRT引擎，需要将其反序列化为可执行的形式。可以使用deserializeCudaEngine函数从序列化的引擎数据中生成TensorRT引擎。

3. 原版build.cu

#include "NvInfer.h"
#include "NvOnnxParser.h" // onnxparser
#include "logger.h"
#include "common.h"
#include "buffers.h"
#include "cassert"

/*
1. Create builder
2. Create Network

*/

int main(int argc, char **argv)
{
    if (argc != 2)  // 命令行参数要等于2 
    {
        std::cerr << "usage:  ./build [onnx_file_path]" << argv[0] << std::endl;
        return -1;
    }

    // onnx_file_path
    const char* onnx_file_path = argv[1];
    // 1. Create Builder
    auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
    // nvinfer::IBuilder *builder = nvinfer1::createInferBuilder(logger);
    if (!builder)
    {
        std::cerr << "create builder failed" << std::endl;
        return -1;
    }

    // 2. Set the input and output names of the network
    auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(1));
    if (!network)
    {
        std::cerr << "create network failed" << std::endl;
        return -1;
    }

    // 3. Parse Onnx configuration
    auto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
    auto parsed = parser->parseFromFile(onnx_file_path, static_cast<int>(sample::gLogger.getReportableSeverity()));
    if (!parsed)
    {
        std::cerr << "parse onnx file failed" << std::endl;
        return -1;
    }

    // 4. Set Image input size 
    // This program only have one input which is one Image at once (1, 3, 640, 640)
    auto input = network->getInput(0);
    auto profile = builder->createOptimizationProfile();
    // set KMIN, KMAX, KOPT
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4{1, 3, 640, 640});
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4{1, 3, 640, 640});
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4{1, 3, 640, 640});

    // 5. Build Network config
    auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    if (!config)
    {
        std::cerr << "create IBuilderConfig failed" << std::endl;
        return -1;
    }

    // config->setFlag() setting precision
    config->setFlag(nvinfer1::BuilderFlag::kFP16);
    // set max batch size
    builder->setMaxBatchSize(1); // infer one image once
    // set max workspace
    config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 30); // 2^30 = 1GB
    
    // Set profileStream
    auto profileStream = samplesCommon::makeCudaStream();
    if (!profileStream)
    {
        std::cerr << "Create CUDA stream failed" << std::endl;
        return -1;
    }
    config->setProfileStream(*profileStream);

    // 6. runing plan
    auto plan = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));
    if (!plan)
    {
        std::cerr << "build engine failed" << std::endl;
        return -1;
    }

    // serialized engine
    std::ofstream engine_file("./weights/my_Engine", std::ios::binary);
    assert(engine_file.is_open() && "Failed to open Engine file");
    engine_file.write((char *)plan->data(), plan->size());

    std::cout<<"Building Engine Successfully"<< 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

4. 解析build.cu的main函数

4.1 命令行参数调整

int argc, char **argv是为了从命令行获取输入参数，以便程序可以根据这些参数进行相应的操作。

argc 是命令行的参数
argv 是指向字符串数组的指针

argv 数组的第一个元素 argv[0] 是程序的名称，而其他元素 argv[1]、argv[2]、……、argv[argc-1] 则是传递给程序的参数。

这里因为在命令行执行可执行文件的时候需要附上onnx的路径, 所以这样子搞一下

if (argc != 2)  // 命令行参数要等于2 
{
    std::cerr << "usage:  ./build [onnx_file_path]" << argv[0] << std::endl;
    return -1;
}

// onnx_file_path
const char* onnx_file_path = argv[1];

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

4.2 创建Builder

第一步都是创建builder,这里使用了TRT自带的全局的Logger, 来自logger.h头文件, 所有使用它的地方都会输出到同一个Logger中。

使用智能指针不用手动去释放了

// ================1. 创建Builder========================
    auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
    /*
    不带智能指针版本:
    nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(logger);
    */
    if (!builder)
    {
        std::cerr << "创建Builder失败" << std::endl;
        return -1;
    }

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

4.3 创建network: builder->Network

显式地创建batch size为1的网络

batch size指定的是每次推理时的输入数据量，一个batch包含的数据量越多，GPU并行处理的能力就越强，同时可以减少数据拷贝的次数，提高推理效率。但是，在某些情况下，例如对于视频数据，我们可能需要单独对每一帧进行推理，因此batch size为1可能更适合。

将ONNX文件的配置解析成TensorRT中的网络结构。通过将网络和解析器一起传递给解析器，可以检查ONNX文件是否正确解析，并将其转换为TensorRT网络，以便可以继续对其进行优化和部署。如果ONNX文件无法解析，后续步骤也就没有意义了。

// 2. Set the input and output names of the network
auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(1));
if (!network)
{
    std::cerr << "create network failed" << std::endl;
    return -1;
}

// 3. Parse Onnx configuration
auto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
auto parsed = parser->parseFromFile(onnx_file_path, static_cast<int>(sample::gLogger.getReportableSeverity()));
if (!parsed)
{
    std::cerr << "parse onnx file failed" << std::endl;
    return -1;
}

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

4.4 network中设置网络的输入

这段代码设置输入尺寸, 要知道这个engine只有一个输入, 就是Image。

getInput(0)拿到第1个输入节点名字(char *), 然后再创建节点profile, 用profile去设置Image尺寸

profile是一种优化配置, 这里用profile设置最小、最大和优化尺寸，可以让TensorRT针对这些特定的尺寸进行优化，以获得更好的性能和效率。这些优化配置可以在创建TensorRT引擎时使用，以便在推理过程中使用最优的配置。

除了控制输入尺寸, TensorRT还有其他的各种手段来优化提高网络的性能, 内存管理, 层次化排序, 算法优化， 量化等等

// 4. Set Image input size 
// This program only have one input which is one Image at once (1, 3, 640, 640)
auto input = network->getInput(0);
auto profile = builder->createOptimizationProfile();
// set KMIN, KMAX, KOPT
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4{1, 3, 640, 640});
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4{1, 3, 640, 640});
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4{1, 3, 640, 640});

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

4.5 创建config: builder->config

创建config, 然后通过config设置精度是FP16，如果不设置的话他默认是FP32，设置为INT8需要额外设置cailbrator

通过再设置最大的workspace, 1<<30的意思就是2^30 = 1GB

通过builder设置最大的batch size

使用makeCudaStream去创建一个profile流, 用于计算每个层之间的运算时间, 并且执行引擎的时候动态调整每层的大小, 获得更好的性能，这里需要跟config关联

auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    if (!config)
    {
        std::cerr << "create IBuilderConfig failed" << std::endl;
        return -1;
    }

    // config->setFlag() setting precision
    config->setFlag(nvinfer1::BuilderFlag::kFP16);
    // set max batch size
    builder->setMaxBatchSize(1); // infer one image once
    // set max workspace
    config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 30); // 2^30 = 1GB
    
    // Set profileStream
    auto profileStream = samplesCommon::makeCudaStream();
    if (!profileStream)
    {
        std::cerr << "Create CUDA stream failed" << std::endl;
        return -1;
    }
    config->setProfileStream(*profileStream);

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

4.6 创建Engine并序列化保存Engine

engine的创建分为两个阶段，分别是构建阶段和序列化阶段。在构建阶段中，会根据输入的网络结构和优化配置创建一个engine对象。而在序列化阶段，将engine对象序列化为一块二进制数据并保存到硬盘上，以便在实际运行中反序列化得到engine对象并使用。

在构建阶段中，由于一些原因（比如内存资源限制），并不是每次都能成功创建一个engine对象。而当创建成功时，可以将这个engine对象看做是一个“计划”（plan），代表了使用TensorRT对网络进行优化之后的运行计划。因此，在代码中通常将engine对象命名为“plan”，以表示这个对象代表了一个运行计划。

// 6. runing plan
auto plan = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));
if (!plan)
{
    std::cerr << "build engine failed" << std::endl;
    return -1;
}

// serialized engine
std::ofstream engine_file("./weights/my_Engine", std::ios::binary);
assert(engine_file.is_open() && "Failed to open Engine file");
engine_file.write((char *)plan->data(), plan->size());

std::cout<<"Building Engine Successfully"<< std::endl;
return 0;

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

5. 原版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"

// 加载模型文件
std::vector<unsigned char> load_engine_file(const std::string &file_name)
{
    std::vector<unsigned char> engine_data;
    std::ifstream engine_file(file_name, std::ios::binary);
    assert(engine_file.is_open() && "Unable to load engine file.");
    engine_file.seekg(0, engine_file.end);
    int length = engine_file.tellg();
    engine_data.resize(length);
    engine_file.seekg(0, engine_file.beg);
    engine_file.read(reinterpret_cast<char *>(engine_data.data()), length);
    return engine_data;
}

int main(int argc, char **argv)
{
    if (argc < 3)
    {
        std::cerr << "用法: " << argv[0] << " <engine_file> <input_path_path>" << std::endl;
        return -1;
    }

    auto engine_file = argv[1];      // 模型文件
    auto input_video_path = argv[2]; // 输入视频文件

    // ========= 1. 创建推理运行时runtime =========
    auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(sample::gLogger.getTRTLogger()));
    if (!runtime)
    {
        std::cout << "runtime create failed" << std::endl;
        return -1;
    }
    // ======== 2. 反序列化生成engine =========
    // 加载模型文件
    auto plan = load_engine_file(engine_file);
    // 反序列化生成engine
    auto mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(plan.data(), plan.size()));
    if (!mEngine)
    {
        return -1;
    }

    // ======== 3. 创建执行上下文context =========
    auto context = std::unique_ptr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
    if (!context)
    {
        std::cout << "context create failed" << std::endl;
        return -1;
    }

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

    auto cap = cv::VideoCapture(input_video_path);

    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/record.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);

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

        cap >> frame;
        if (frame.empty())
        {
            std::cout << "文件处理完毕" << std::endl;
            break;
        }
        frame_index++;

        // 输入预处理（实现了对输入图像处理的gpu 加速)
        process_input(frame, (float *)buffers.getDeviceBuffer(kInputTensorName));
        // ========== 5. 执行推理 =========
        context->executeV2(buffers.getDeviceBindings().data());
        // 拷贝回host
        buffers.copyOutputToHost();

        // 从buffer manager中获取模型输出
        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);     // 检测到的目标框
        // 执行nms（非极大值抑制），得到最后的检测框
        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 j = 0; j < bboxs.size(); j++)
        {
            cv::Rect r = get_rect(frame, bboxs[j].bbox);
            cv::rectangle(frame, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
            cv::putText(frame, std::to_string((int)bboxs[j].class_id), cv::Point(r.x, r.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);

        // cv::imshow("frame", frame);
        // 写入视频文件
        writer.write(frame);
        std::cout << "处理完第" << frame_index << "帧" << std::endl;
        if (cv::waitKey(1) == 27)
            break;
    }
    // ========== 6. 释放资源 =========
    // 因为使用了unique_ptr，所以不需要手动释放

    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

6. 解析版runtime.cu

6.1 加载Engine

典型的读取二进制文件

static_cast 是C++中的一种类型转换操作，用于在编译时执行类型转换。对于某些类型，例如 int 和 float，static_cast 可以执行从一种类型到另一种类型的转换。但是，对于指针类型，static_cast 只允许在安全的情况下进行转换。

对于指向无符号字符类型的指针 unsigned char * 和指向字符类型的指针 char *，它们在内存表示和语义上是相同的，因此可以通过 static_cast 在这两种类型之间进行转换。但是，在标准C++语言中， static_cast 无法将 unsigned char * 类型直接转换为 char * 类型，因为它们之间没有隐式转换的关系。因此，在这种情况下，需要使用 reinterpret_cast 执行指针类型转换。

// 加载模型文件
std::vector<unsigned char> load_engine_file(const std::string &file_name)
{
    std::vector<unsigned char> engine_data;
    std::ifstream engine_file(file_name, std::ios::binary);
    assert(engine_file.is_open() && "Unable to load engine file.");
    engine_file.seekg(0, engine_file.end);
    int length = engine_file.tellg();
    engine_data.resize(length);
    engine_file.seekg(0, engine_file.beg);
    engine_file.read(reinterpret_cast<char *>(engine_data.data()), length);
    return engine_data;
}

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

6.2 保证输入是三个，因为runtime的可执行文件要加上engine path 和 推理文件 path

if (argc < 3)
    {
        std::cerr << "用法: " << argv[0] << " <engine_file> <input_path_path>" << std::endl;
        return -1;
    }

    auto engine_file = argv[1];      // 模型文件
    auto input_video_path = argv[2]; // 输入视频文件

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

6.3 创建推理运行时runtime

auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(sample::gLogger.getTRTLogger()));
if (!runtime)
{
    std::cout << "runtime create failed" << std::endl;
    return -1;
}

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

6.4 反序列化生成engine

这里的engine可能会被多次引用所以用std::shared_ptr

// ======== 2. 反序列化生成engine =========
// 加载模型文件
auto plan = load_engine_file(engine_file);
// 反序列化生成engine
auto mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(plan.data(), plan.size()));
if (!mEngine)
{
    return -1;
}

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

6.5 创建上下文context

简单点说，context是从engine获得推理任务的定义，然后通过runtime这个运行是的环境管理设备，进行推理任务

// ======== 3. 创建执行上下文context =========
auto context = std::unique_ptr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
if (!context)
{
    std::cout << "context create failed" << std::endl;
    return -1;
}

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

6.6 创建输入输出缓存区

在这里，首先创建了一个 BufferManager 对象 buffers，它负责为模型的输入和输出数据分配和管理缓冲区。 BufferManager 类是由 TensorRT 示例代码提供的一个实用类.

每一步操作都从buffer拿数据搞完了再放回去

buffer是与TensorRT engine相关联的。在TensorRT中，输入和输出数据需要通过BufferManager来分配和管理，其中BufferManager会创建和存储一系列的缓冲区来存储输入和输出Tensor。在运行时，我们需要将输入数据放入输入Tensor的缓冲区中，然后通过执行上下文执行推理，推理结果也会被存储在输出Tensor的缓冲区中。因此，BufferManager是一个用于在运行时分配和管理输入和输出Tensor缓冲区的实用工具。

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

  • 1
  • 2

6.7 拿到图像的信息创建vedio

拿到图像的宽高, 拿来申请内存, 确保有足够的内存处理每一帧图像

创建writer, 写入视频

// 拿到图像信息，宽，高，fps
    auto cap = cv::VideoCapture(vedio_path);

    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/record.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
  • 14
  • 15

6.8 while循环处理每一帧

使用tensorrtx的处理

预处理->执行推理->拷贝回host->后处理(NMS)->恢复bboxs

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

        cap >> frame;
        if (frame.empty())
        {
            std::cout << "文件处理完成" << std::endl;
            break;
        }
        frame_index++;

        // 输入预处理: 把buffers(缓存区的数据)放到GPU上进行预处理操作
        process_input(frame, (float *)buffers.getDeviceBuffer(kInputTensorName));

        // 5. 执行推理
        context->executeV2(buffers.getDeviceBindings().data());

        // 拷贝回host
        buffers.copyOutputToHost();

        /*
        从buffer manager中获取模型输出, 检测数量, 类别, 置信度, bounding boxes
        getHostBuffer() 是一个 (void *) 类型, 需要转换成对应类型的*
        */
        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);

        // 执行nms
        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 j = 0; j < bboxs.size(); j++)
        {
            cv::Rect r = get_rect(frame, bboxs[j].bbox);
            cv::rectangle(frame, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
            cv::putText(frame, std::to_string((int)bboxs[j].class_id), cv::Point(r.x, r.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;

        // 不用释放资源
    }
    

  • 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
声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/我家小花儿/article/detail/261428
推荐阅读
相关标签

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

  

闽ICP备14008679号