> 文章列表 > TensorRT创建Engine并推理engine

TensorRT创建Engine并推理engine

TensorRT创建Engine并推理engine

1. 验证集数据集

               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|██████████| 84/all       1000      28423      0.451      0.374      0.376      0.209pedestrians       1000      17833      0.737      0.855       0.88      0.609riders       1000        185      0.545      0.492      0.521      0.256
partially-visible-person       1000       9335      0.456      0.338      0.336      0.125ignore-regions       1000        409       0.37      0.138      0.121     0.0485crowd       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)

2. 视频流推理

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

计算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上的推理是

TensorRT创建Engine并推理engine

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_pathconst char* onnx_file_path = argv[1];// 1. Create Builderauto 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 networkauto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(1));if (!network){std::cerr << "create network failed" << std::endl;return -1;}// 3. Parse Onnx configurationauto 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, KOPTprofile->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 configauto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());if (!config){std::cerr << "create IBuilderConfig failed" << std::endl;return -1;}// config->setFlag() setting precisionconfig->setFlag(nvinfer1::BuilderFlag::kFP16);// set max batch sizebuilder->setMaxBatchSize(1); // infer one image once// set max workspaceconfig->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 30); // 2^30 = 1GB// Set profileStreamauto profileStream = samplesCommon::makeCudaStream();if (!profileStream){std::cerr << "Create CUDA stream failed" << std::endl;return -1;}config->setProfileStream(*profileStream);// 6. runing planauto plan = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));if (!plan){std::cerr << "build engine failed" << std::endl;return -1;}// serialized enginestd::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;
}

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

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

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

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

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 precisionconfig->setFlag(nvinfer1::BuilderFlag::kFP16);// set max batch sizebuilder->setMaxBatchSize(1); // infer one image once// set max workspaceconfig->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 30); // 2^30 = 1GB// Set profileStreamauto profileStream = samplesCommon::makeCudaStream();if (!profileStream){std::cerr << "Create CUDA stream failed" << std::endl;return -1;}config->setProfileStream(*profileStream);

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;

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);// 反序列化生成engineauto 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());// 拷贝回hostbuffers.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;
}

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

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]; // 输入视频文件

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

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

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

6.6 创建输入输出缓存区

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

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

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

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

6.7 拿到图像的信息创建vedio

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

创建writer, 写入视频

// 拿到图像信息,宽,高,fpsauto 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);

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());// 拷贝回hostbuffers.copyOutputToHost();/*从buffer manager中获取模型输出, 检测数量, 类别, 置信度, bounding boxesgetHostBuffer() 是一个 (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);// 执行nmsstd::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;// 不用释放资源}