TRT部署ONNX

 也算是基础知识了 就搬来了 还是要说大佬勿怪   其实应该算是 转成trt默认 使用 这样说确切

不过我一般不这么用yolo直接 用王鑫宇 大佬工具转wts然后载转trt模型 大伙去git上找把 都支持到v7了 是阿拉伯联合酋长国阿布拉卡那版哦

TensorRT是英伟达官方提供的一个高性能深度学习推理优化库，支持C++和Python两种编程语言API。通常情况下深度学习模型部署都会追求效率，尤其是在嵌入式平台上，所以一般会选择使用C++来做部署。

这里将以YOLOv5为例详细介绍如何使用TensorRT的C++版本API来部署ONNX模型，使用的TensorRT版本为8.4.1.5，如果使用其他版本可能会存在某些函数与本文描述的不一致。另外，使用TensorRT 7会导致YOLOv5的输出结果与期望不一致，请注意。

导出ONNX模型

YOLOv5使用PyTorch框架进行训练，可以使用官方代码仓库中的export.py脚本把PyTorch模型转换为ONNX模型:

python export.py --weights yolov5x.pt --include onnx --imgsz 640 640  

准备模型输入数据

如果想用YOLOv5对图像做目标检测，在将图像输入给模型之前还需要做一定的预处理操作，预处理操作应该与模型训练时所做的操作一致。YOLOv5的输入是RGB格式的3通道图像，图像的每个像素需要除以255来做归一化，并且数据要按照CHW的顺序进行排布。所以YOLOv5的预处理大致可以分为两个步骤：

1. 将原始输入图像缩放到模型需要的尺寸，比如640x640。这一步需要注意的是，原始图像是按照等比例进行缩放的，如果缩放后的图像某个维度上比目标值小，那么就需要进行填充。举个例子：假设输入图像尺寸为768x576，模型输入尺寸为640x640，按照等比例缩放的原则缩放后的图像尺寸为640x480，那么在y方向上还需要填充640-480=160（分别在图像的顶部和底部各填充80）。来看一下实现代码：

  

   

    

     

    
    

     

      cv::Mat input_image = cv::
      imread(
      "dog.jpg");  
     
    
   

    

     

    
    

     

      cv::Mat resize_image;  
     
    
   

    

     

    
    

     

      const 
      int model_width = 
      640;  
     
    
   

    

     

    
    

     

      const 
      int model_height = 
      640;  
     
    
   

    

     

    
    

     

      const 
      float ratio = std::
      min(model_width / (input_image.cols * 
      1.0f),  
     
    
   

    

     

    
    

     

                                    model_height / (input_image.rows * 
      1.0f));  
     
    
   

    

     

    
    

     

      // 等比例缩放  
     
    
   

    

     

    
    

     

      const 
      int border_width = input_image.cols * ratio;  
     
    
   

    

     

    
    

     

      const 
      int border_height = input_image.rows * ratio;  
     
    
   

    

     

    
    

     

      // 计算偏移值  
     
    
   

    

     

    
    

     

      const 
      int x_offset = (model_width - border_width) / 
      2;  
     
    
   

    

     

    
    

     

      const 
      int y_offset = (model_height - border_height) / 
      2;  
     
    
   

    

     

    
    

     

      cv::
      resize(input_image, resize_image, cv::
      Size(border_width, border_height));  
     
    
   

    

     

    
    

     

      cv::
      copyMakeBorder(resize_image, resize_image, y_offset, y_offset, x_offset,  
     
    
   

    

     

    
    

     

                          x_offset, cv::BORDER_CONSTANT, cv::
      Scalar(
      114, 
      114, 
      114));  
     
    
   

    

     

    
    

     

      // 转换为RGB格式  
     
    
   

    

     

    
    

     

      cv::
      cvtColor(resize_image, resize_image, cv::COLOR_BGR2RGB);  
     
    
  
 

  
 

图像这样处理后的效果如下图所示，顶部和底部的灰色部分是填充后的效果。

  2. 对图像像素做归一化操作，并按照CHW的顺序进行排布。这一步的操作比较简单，直接看代码吧：

  

   

    

     

    
    

     

      input_blob = 
      new 
      float[model_height * model_width * 
      3];  
     
    
   

    

     

    
    

     

      const 
      int channels = resize_image.
      channels();  
     
    
   

    

     

    
    

     

      const 
      int width = resize_image.cols;  
     
    
   

    

     

    
    

     

      const 
      int height = resize_image.rows;  
     
    
   

    

     

    
    

     

      for (
      int c = 
      0; c < channels; c++) {  
     
    
   

    

     

    
    

     

        
      for (
      int h = 
      0; h < height; h++) {  
     
    
   

    

     

    
    

     

          
      for (
      int w = 
      0; w < width; w++) {  
     
    
   

    

     

    
    

     

            input_blob[c * width * height + h * width + w] =  
     
    
   

    

     

    
    

     

                resize_image.
      at<cv::Vec3b>(h, w)[c] / 
      255.0f;  
     
    
   

    

     

    
    

     

          }  
     
    
   

    

     

    
    

     

        }  
     
    
   

    

     

    
    

     

      }  
     
    
  

ONNX模型部署

1. 模型优化与序列化

要使用TensorRT的C++ API来部署模型，首先需要包含头文件NvInfer.h。

#include "NvInfer.h"  

TensorRT所有的编程接口都被放在命名空间nvinfer1中，并且都以字母I为前缀，比如ILogger、IBuilder等。使用TensorRT部署模型首先需要创建一个IBuilder对象，创建之前还要先实例化ILogger接口：

  

   

    

     

    
    

     

      class 
      MyLogger : 
      public nvinfer1::ILogger {  
     
    
   

    

     

    
    

     

       
      public:  
     
    
   

    

     

    
    

     

        
      explicit MyLogger(nvinfer1::ILogger::Severity severity =  
     
    
   

    

     

    
    

     

                              nvinfer1::ILogger::Severity::kWARNING)  
     
    
   

    

     

    
    

     

            : severity_(severity) {}  
     
    
   

    

     

    
    

     

        
     
    
   

    

     

    
    

     

        
      void log(nvinfer1::ILogger::Severity severity,  
     
    
   

    

     

    
    

     

                 
      const 
      char *msg) 
      noexcept 
      override {  
     
    
   

    

     

    
    

     

          
      if (severity <= severity_) {  
     
    
   

    

     

    
    

     

            std::cerr << msg << std::endl;  
     
    
   

    

     

    
    

     

          }  
     
    
   

    

     

    
    

     

        }  
     
    
   

    

     

    
    

     

        nvinfer1::ILogger::Severity severity_;  
     
    
   

    

     

    
    

     

      };  
     
    
  

上面的代码默认会捕获级别大于等于WARNING的日志信息并在终端输出。实例化ILogger接口后，就可以创建IBuilder对象：

  

   

    

     

    
    

     

      MyLogger logger;  
     
    
   

    

     

    
    

     

      nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(logger);  
     
    
  

创建IBuilder对象后，优化一个模型的第一步是要构建模型的网络结构。

  

   

    

     

    
    

     

      const uint32_t explicit_batch = 
      1U << static_cast<uint32_t>(  
     
    
   

    

     

    
    

     

                nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);  
     
    
   

    

     

    
    

     

      nvinfer1::INetworkDefinition *network = builder
      ->
      createNetworkV2(explicit_batch);  
     
    
  

模型的网络结构有两种构建方式，一种是使用TensorRT的API一层一层地去搭建，这种方式比较麻烦；另外一种是直接从ONNX模型中解析出模型的网络结构，这需要ONNX解析器来完成。由于我们已经有现成的ONNX模型了，所以选择第二种方式。TensorRT的ONNX解析器接口被封装在头文件NvOnnxParser.h中，命名空间为nvonnxparser。创建ONNX解析器对象并加载模型的代码如下：

  

   

    

     

    
    

     

      const std::string onnx_model = 
      "yolov5m.onnx";  
     
    
   

    

     

    
    

     

      nvonnxparser::IParser *parser = nvonnxparser::
      createParser(*network, logger);  
     
    
   

    

     

    
    

     

      parser
      ->
      parseFromFile(model_path.
      c_str(),  
     
    
   

    

     

    
    

     

          static_cast<int>(nvinfer1::ILogger::Severity::kERROR))  
     
    
   

    

     

    
    

     

      // 如果有错误则输出错误信息  
     
    
   

    

     

    
    

     

      for (int32_t i = 
      0; i < parser
      ->
      getNbErrors(); ++i) {  
     
    
   

    

     

    
    

     

          std::cout << parser
      ->
      getError(i)
      ->
      desc() << std::endl;  
     
    
   

    

     

    
    

     

      }  
     
    
  

模型解析成功后，需要创建一个IBuilderConfig对象来告诉TensorRT该如何对模型进行优化。这个接口定义了很多属性，其中最重要的一个属性是工作空间的最大容量。在网络层实现过程中通常会需要一些临时的工作空间，这个属性会限制最大能申请的工作空间的容量，如果容量不够的话会导致该网络层不能成功实现而导致错误。另外，还可以通过这个对象设置模型的数据精度。TensorRT默认的数据精度为FP32，我们还可以设置FP16或者INT8，前提是该硬件平台支持这种数据精度。

  

   

    

     

    
    

     

      nvinfer1::IBuilderConfig *config = builder
      ->
      createBuilderConfig();  
     
    
   

    

     

    
    

     

      config
      ->
      setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 
      1U << 
      25);  
     
    
   

    

     

    
    

     

      if (builder
      ->
      platformHasFastFp16()) {  
     
    
   

    

     

    
    

     

        config
      ->
      setFlag(nvinfer1::BuilderFlag::kFP16);  
     
    
   

    

     

    
    

     

      }  
     
    
  

设置IBuilderConfig属性后，就可以启动优化引擎对模型进行优化了，这个过程需要一定的时间，在嵌入式平台上可能会比较久一点。经过TensorRT优化后的序列化模型被保存到IHostMemory对象中，我们可以将其保存到磁盘中，下次使用时直接加载这个经过优化的模型即可，这样就可以省去漫长的等待模型优化的过程。我一般习惯把序列化模型保存到一个后缀为.engine的文件中。

  

   

    

     

    
    

     

      nvinfer1::IHostMemory *serialized_model =  
     
    
   

    

     

    
    

     

            builder->
      buildSerializedNetwork(*network, *config);  
     
    
   

    

     

    
    

     

        
     
    
   

    

     

    
    

     

      // 将模型序列化到engine文件中  
     
    
   

    

     

    
    

     

      std::stringstream engine_file_stream;  
     
    
   

    

     

    
    

     

      engine_file_stream
      .seekg(
      0, engine_file_stream.beg);  
     
    
   

    

     

    
    

     

      engine_file_stream
      .write(static_cast<const char *>(serialized_model->data()),  
     
    
   

    

     

    
    

     

                              serialized_model->
      size());  
     
    
   

    

     

    
    

     

      const std::string engine_file_path = 
      "yolov5m.engine";  
     
    
   

    

     

    
    

     

      std::ofstream 
      out_file(engine_file_path);  
     
    
   

    

     

    
    

     

      assert(out_file.is_open());  
     
    
   

    

     

    
    

     

      out_file << engine_file_stream
      .rdbuf();  
     
    
   

    

     

    
    

     

      out_file
      .close();  
     
    
  

由于IHostMemory对象保存了模型所有的信息，所以前面创建的IBuilder、IParser等对象已经不再需要了，可以通过delete进行释放。

  

   

    

     

    
    

     

      delete config;  
     
    
   

    

     

    
    

     

      delete parser;  
     
    
   

    

     

    
    

     

      delete network;  
     
    
   

    

     

    
    

     

      delete builder;  
     
    
  

IHostMemory对象用完后也可以通过delete进行释放。

2. 模型反序列化

通过上一步得到优化后的序列化模型后，如果要用模型进行推理，那么还需要创建一个IRuntime接口的实例，然后通过其模型反序列化接口去创建一个ICudaEngine对象：

  

   

    

     

    
    

     

      nvinfer1::IRuntime *runtime = nvinfer1::
      createInferRuntime(logger);  
     
    
   

    

     

    
    

     

      nvinfer1::ICudaEngine *engine = runtime->
      deserializeCudaEngine(  
     
    
   

    

     

    
    

     

          serialized_model->
      data(), serialized_model->
      size());  
     
    
   

    

     

    
    

     

        
     
    
   

    

     

    
    

     

      delete serialized_model;  
     
    
   

    

     

    
    

     

      delete runtime;  
     
    
  

如果是直接从磁盘中加载.engine文件也是差不多的步骤，首先从.engine文件中把模型加载到内存中，然后再通过IRuntime接口对模型进行反序列化即可。

  

   

    

     

    
    

     

      const std::string engine_file_path = 
      "yolov5m.engine";  
     
    
   

    

     

    
    

     

      std::stringstream engine_file_stream;  
     
    
   

    

     

    
    

     

      engine_file_stream.
      seekg(
      0, engine_file_stream.beg);  
     
    
   

    

     

    
    

     

      std::ifstream 
      ifs(engine_file_path);  
     
    
   

    

     

    
    

     

      engine_file_stream << ifs.
      rdbuf();  
     
    
   

    

     

    
    

     

      ifs.
      close();  
     
    
   

    

     

    
    

     

        
     
    
   

    

     

    
    

     

      engine_file_stream.
      seekg(
      0, std::ios::end);  
     
    
   

    

     

    
    

     

      const int model_size = engine_file_stream.
      tellg();  
     
    
   

    

     

    
    

     

      engine_file_stream.
      seekg(
      0, std::ios::beg);  
     
    
   

    

     

    
    

     

      void *model_mem = 
      malloc(model_size);  
     
    
   

    

     

    
    

     

      engine_file_stream.
      read(static_cast<
      char *>(model_mem), model_size);  
     
    
   

    

     

    
    

     

        
     
    
   

    

     

    
    

     

      nvinfer1::IRuntime *runtime = nvinfer1::
      createInferRuntime(logger);  
     
    
   

    

     

    
    

     

      nvinfer1::ICudaEngine *engine = runtime
      ->
      deserializeCudaEngine(model_mem, model_size);  
     
    
   

    

     

    
    

     

        
     
    
   

    

     

    
    

     

      delete runtime;  
     
    
   

    

     

    
    

     

      free(model_mem);  
     
    
  
 

  
 

3. 模型推理

ICudaEngine对象中存放着经过TensorRT优化后的模型，不过如果要用模型进行推理则还需要通过createExecutionContext()函数去创建一个IExecutionContext对象来管理推理的过程：

nvinfer1::IExecutionContext *context = engine->createExecutionContext();  

现在让我们先来看一下使用TensorRT框架进行模型推理的完整流程：

1. 对输入图像数据做与模型训练时一样的预处理操作。

2. 把模型的输入数据从CPU拷贝到GPU中。

3. 调用模型推理接口进行推理。

4. 把模型的输出数据从GPU拷贝到CPU中。

5. 对模型的输出结果进行解析，进行必要的后处理后得到最终的结果。

由于模型的推理是在GPU上进行的，所以会存在搬运输入、输出数据的操作，因此有必要在GPU上创建内存区域用于存放输入、输出数据。模型输入、输出的尺寸可以通过ICudaEngine对象的接口来获取，根据这些信息我们可以先为模型分配输入、输出缓存区。

  

   

    

     

    
    

     

      void *buffers[
      2];  
     
    
   

    

     

    
    

     

      // 获取模型输入尺寸并分配GPU内存  
     
    
   

    

     

    
    

     

      nvinfer1::
      Dims input_dim = engine->
      getBindingDimensions(
      0);  
     
    
   

    

     

    
    

     

      int input_size = 
      1;  
     
    
   

    

     

    
    

     

      for
       (
      int j = 
      0; j < input_dim.nbDims; ++j) {  
     
    
   

    

     

    
    

     

        input_size *= input_dim.d[j];  
     
    
   

    

     

    
    

     

      }  
     
    
   

    

     

    
    

     

      cudaMalloc(&buffers[
      0], input_size * 
      sizeof(
      float));  
     
    
   

    

     

    
    

     

        
     
    
   

    

     

    
    

     

      // 获取模型输出尺寸并分配GPU内存  
     
    
   

    

     

    
    

     

      nvinfer1::
      Dims output_dim = engine->
      getBindingDimensions(
      1);  
     
    
   

    

     

    
    

     

      int output_size = 
      1;  
     
    
   

    

     

    
    

     

      for
       (
      int j = 
      0; j < output_dim.nbDims; ++j) {  
     
    
   

    

     

    
    

     

        output_size *= output_dim.d[j];  
     
    
   

    

     

    
    

     

      }  
     
    
   

    

     

    
    

     

      cudaMalloc(&buffers[
      1], output_size * 
      sizeof(
      float));  
     
    
   

    

     

    
    

     

        
     
    
   

    

     

    
    

     

      // 给模型输出数据分配相应的CPU内存  
     
    
   

    

     

    
    

     

      float *output_buffer = 
      new 
      float[output_size]();  
     
    
  
 

  
 

到这一步，如果你的输入数据已经准备好了，那么就可以调用TensorRT的接口进行推理了。通常情况下，我们会调用IExecutionContext对象的enqueueV2()函数进行异步地推理操作，该函数的第二个参数为CUDA流对象，第三个参数为CUDA事件对象，这个事件表示该执行流中输入数据已经使用完，可以挪作他用了。

  

   

    

     

    
    

     

      cudaStream_t stream;  
     
    
   

    

     

    
    

     

      cudaStreamCreate(&stream);  
     
    
   

    

     

    
    

     

      // 拷贝输入数据  
     
    
   

    

     

    
    

     

      cudaMemcpyAsync(buffers[
      0], input_blob,input_size * sizeof(float),  
     
    
   

    

     

    
    

     

                        cudaMemcpyHostToDevice, stream);  
     
    
   

    

     

    
    

     

      // 执行推理  
     
    
   

    

     

    
    

     

      context->
      enqueueV2(buffers, stream, nullptr);  
     
    
   

    

     

    
    

     

      // 拷贝输出数据  
     
    
   

    

     

    
    

     

      cudaMemcpyAsync(output_buffer, buffers[
      1],output_size * sizeof(float),  
     
    
   

    

     

    
    

     

                        cudaMemcpyDeviceToHost, stream);  
     
    
   

    

     

    
    

     

        
     
    
   

    

     

    
    

     

      cudaStreamSynchronize(stream);  
     
    
  

模型推理成功后，其输出数据被拷贝到output_buffer中，接下来我们只需按照YOLOv5的输出数据排布规则去解析即可。

4. 小结

在介绍如何解析YOLOv5输出数据之前，我们先来总结一下用TensorRT框架部署ONNX模型的基本流程。

 

如上图所示，主要步骤如下：                            whaosoft aiot http://143ai.com 

1. 实例化Logger;

2. 创建Builder;

3. 创建Network;

4. 使用Parser解析ONNX模型，构建Network；

5. 设置Config参数；

6. 优化网络，序列化模型；

7. 反序列化模型；

8. 拷贝模型输入数据（HostToDevice），执行模型推理；

9. 拷贝模型输出数据（DeviceToHost），解析结果。

解析模型输出结果

YOLOv5有3个检测头，如果模型输入尺寸为640x640，那么这3个检测头分别在80x80、40x40和20x20的特征图上做检测。让我们先用Netron工具来看一下YOLOv5 ONNX模型的结构，可以看到，YOLOv5的后处理操作已经被包含在模型中了（如下图红色框内所示），3个检测头分支的结果最终被组合成一个张量作为输出。

 yolov5m

YOLOv5的3个检测头一共有(80x80+40x40+20x20)x3=25200个输出单元格，每个单元格输出x,y,w,h,objectness这5项再加80个类别的置信度总共85项内容。经过后处理操作后，目标的坐标值已经被恢复到以640x640为参考的尺寸，如果需要恢复到原始图像尺寸，只需要除以预处理时的缩放因子即可。这里有个问题需要注意：由于在做预处理的时候图像做了填充，原始图像并不是被缩放成640x640而是640x480，使得输入给模型的图像的顶部被填充了一块高度为80的区域，所以在恢复到原始尺寸之前，需要把目标的y坐标减去偏移量80。

详细的解析代码如下：

  

   

    

     

    
    

     

      float *ptr = output_buffer;  
     
    
   

    

     

    
    

     

      for (
      int i = 
      0; i < 
      25200; ++i) {  
     
    
   

    

     

    
    

     

        
      const 
      float objectness = ptr[
      4];  
     
    
   

    

     

    
    

     

        
      if (objectness >= 
      0.45f) {  
     
    
   

    

     

    
    

     

          
      const 
      int label =  
     
    
   

    

     

    
    

     

              std::
      max_element(ptr + 
      5, ptr + 
      85) - (ptr + 
      5);  
     
    
   

    

     

    
    

     

          
      const 
      float confidence = ptr[
      5 + label] * objectness;  
     
    
   

    

     

    
    

     

          
      if (confidence >= 
      0.25f) {  
     
    
   

    

     

    
    

     

            
      const 
      float bx = ptr[
      0];  
     
    
   

    

     

    
    

     

            
      const 
      float by = ptr[
      1];  
     
    
   

    

     

    
    

     

            
      const 
      float bw = ptr[
      2];  
     
    
   

    

     

    
    

     

            
      const 
      float bh = ptr[
      3];  
     
    
   

    

     

    
    

     

        
     
    
   

    

     

    
    

     

            Object obj;  
     
    
   

    

     

    
    

     

            
      // 这里要减掉偏移值  
     
    
   

    

     

    
    

     

            obj.box.x = (bx - bw * 
      0.5f - x_offset) / ratio;  
     
    
   

    

     

    
    

     

            obj.box.y = (by - bh * 
      0.5f - y_offset) / ratio;  
     
    
   

    

     

    
    

     

            obj.box.width = bw / ratio;  
     
    
   

    

     

    
    

     

            obj.box.height = bh / ratio;  
     
    
   

    

     

    
    

     

            obj.label = label;  
     
    
   

    

     

    
    

     

            obj.confidence = confidence;  
     
    
   

    

     

    
    

     

            objs->
      push_back(std::
      move(obj));  
     
    
   

    

     

    
    

     

          }  
     
    
   

    

     

    
    

     

        }  
     
    
   

    

     

    
    

     

        ptr += 
      85;  
     
    
   

    

     

    
    

     

      }  
      // i loop  
     
    
  
 

  
 

对解析出的目标做非极大值抑制（NMS）操作后，检测结果如下图所示：

 

本文以YOLOv5为例通过大量的代码一步步讲解如何使用TensorRT框架部署ONNX模型，主要目的是希望读者能够通过本文学习到TensorRT模型部署的基本流程，比如如何准备输入数据、如何调用API用模型做推理、如何解析模型的输出结果。如何部署YOLOv5模型并不是本文的重点，重点是要掌握使用TensorRT部署ONNX模型的基本方法，这样才会有举一反三的效果。

 

转载：https://blog.csdn.net/qq_29788741/article/details/127542614