Yolo系列模型使用TensorRT-C++推理学习

sansheng0208

已于 2023-12-11 15:20:57 修改

阅读量1.1k

点赞数 18

分类专栏： c++学习文章标签： YOLO c++ 学习计算机视觉深度学习

于 2023-12-06 18:47:20 首次发布

本文链接：https://blog.csdn.net/sansheng0208/article/details/134811017

版权

c++学习专栏收录该内容

1 篇文章 0 订阅

订阅专栏

前言

最近开始机器视觉、深度学习在Linux c++方向（大概就是这个方向）的学习，其实主要是为了学习c++在Linux环境做一些工程，刚好自己是做机器视觉相关算法的，但之前大部分用的都是python，以及jetson盒子相关配套的东西，所以就从TensorRT-C++推理开始学习了。
这里也是记录一下自己的学习过程，因为c++较为繁琐复杂，通过笔记来记录加深印象。

参考

首先贴出学习代码的来源：
TensorRT-For-YOLO-Series 后来发现该代码应该是参照YoloX写的
YoloX-TRT
一些环境的安装，譬如我是用vscode在linux远程服务器建立docker 容器编写、编译、运行c++代码的，这些也都参考了很多回答和文章，如果有需要可以评论我专门写一篇文章用于介绍前面的准备工作。

开始学习！

一切都从环境安装开始，所幸本次学习的代码环境安装并不复杂。

环境安装

机器配置就是普通的3090两卡服务器，采用的是docker镜像做容器的方式配置的环境。
docker镜像： nvcr.io/nvidia/tensorrt:23.10-py3 （docker镜像版本里面的cuda版本要能与宿主机的cuda driver适应）
起容器命令： sudo docker run -it -v /data/smh:/workdir --workdir=/workdir/ --network=host --runtime nvidia --gpus all --shm-size 16g --name trt nvcr.io/nvidia/tensorrt:23.10-py3 /bin/bash
进去之后看一下cmake版本，cmake --version，我这里是3.24.0。
OpenCV安装 ：可以选择自己安装编译opencv的包（这个需要自己多查查，我也是查了好久结合好几个回答，才安装编译成功），但是这样的话需要修改工程里面cmakelist相关路径，推荐不太懂cmake的伙伴直接通过系统安装。
可能需要安装的依赖（g++，python库，读取视频流的库等，可以先试试安装opencv的命令，出问题再回来执行安装这些库）：

apt update && apt upgrade
apt-get install build-essential libgtk2.0-dev libgtk-3-dev libavcodec-dev libavformat-dev libjpeg-dev libswscale-dev libtiff5-dev

apt install python3-dev python3-numpy
apt install libgstreamer-plugins-base1.0-dev libgstreamer1.0-dev
apt install libpng-dev libopenexr-dev libtiff-dev libwebp-dev

系统安装OpenCV的命令：apt-get update -y && apt-get install -y libopencv-dev
至此，TensorRT相关的包，拉取英伟达的镜像里面基本都安装好了，大概率需要自己安装opencv，其余的应该就没什么了，所以环境搭建基本完成。

代码分析

TensorRT-For-YOLO-Series/cpp/end2end/main.cpp
首先是类的一些基本定义，这里一共两个类 Logger和yolo。
Logger类继承自nvinfer1::ILogger，通过实现ILogger接口，并将其传递给TensorRT相关的API，可以捕获和处理TensorRT运行时产生的日志消息。
这个ILogger类是TensorRT建立、使用各种api的基础。

// Logger类，继承自父类  nvinfer1::ILogger
class Logger : public ILogger {
 public:
 // noexcept override 表示不允许被重载
  void log(Severity severity, const char* msg) noexcept override {
    if (severity != Severity::kINFO) {
      std::cout << msg << std::endl;
    }
  }
};

yolo类是针对于yolo模型定义的类，里面主要包含一些对于图片的预处理，加载trt模型、推理、销毁内存等。

class Yolo {
 public:
  Yolo(char* model_path); // 构造函数，初始化空间、cuda流等等
  // letterbox用于resize图片并填充图片
  float letterbox(
      const cv::Mat& image,
      cv::Mat& out_image,
      const cv::Size& new_shape,
      int stride,
      const cv::Scalar& color,
      bool fixed_shape,
      bool scale_up);
  float* blobFromImage(cv::Mat& img); //相当于标准化，并且将图片转化为一个float数组
  void draw_objects(const cv::Mat& img, float* Boxes, int* ClassIndexs, int* BboxNum);
  void Init(char* model_path);
  void Infer(
      int aWidth,
      int aHeight,
      int aChannel,
      unsigned char* aBytes,
      float* Boxes,
      int* ClassIndexs,
      int* BboxNum);
  ~Yolo();

 private:
  nvinfer1::ICudaEngine* engine = nullptr;
  nvinfer1::IRuntime* runtime = nullptr;
  nvinfer1::IExecutionContext* context = nullptr;
  cudaStream_t stream = nullptr; 
  void* buffs[5]; // 为输入输出开辟的内存空间 
  int iH, iW, in_size, out_size1, out_size2, out_size3, out_size4;
  Logger gLogger; // 要创建Builder、Runtime等类的实例，必须先初始化ILogger实例
};

赋值构造函数只需要一个参数：模型地址，加载模型并反序列化到buffer，再初始化一些cuda流等等。

根据路径读取文件流，并计算文件的size
根据gLogger创建一个Runtime实例
调用initLibNvInferPlugins 注册所有plugin（这里用了nms的plugin）
反序列化模型，把模型文件加载进buffer
初始化一个context
通过model里面的binding信息，得到输入输出的size，根据size申请cuda的buffer
初始化一个cuda stream

Yolo::Yolo(char* model_path) {
  ifstream ifile(model_path, ios::in | ios::binary); // 读取的文件流
  if (!ifile) {
    cout << "read serialized file failed\n";
    std::abort(); // 直接发送异常，终止程序
  }

  ifile.seekg(0, ios::end); // seekg设置指针位置，这里设置到文件末尾
  const int mdsize = ifile.tellg(); // tellg是当前指针位置，这里计算出文件的size
  // 多次操作同一个文件流，为了避免后面的操作因为文件流处于eof状态而无法正常进行，在每种操作方法前，都使用clear()方法将文件流状态重设为默认的“goodbit”状态。
  ifile.clear();  
  ifile.seekg(0, ios::beg);
  vector<char> buf(mdsize);
  ifile.read(&buf[0], mdsize); // 把文件流读取到buf
  ifile.close();
  cout << "model size: " << mdsize << endl;

  runtime = nvinfer1::createInferRuntime(gLogger); //创建一个Runtime实例
  initLibNvInferPlugins(&gLogger, ""); //调用initLibNvInferPlugins 注册所有插件
  engine = runtime->deserializeCudaEngine((void*)&buf[0], mdsize);   // 反序列化模型，入参是加载模型的buffer和模型的size
  // getBindingIndex通过tensor的name获取tensor的序号，再传入getBindingDimensions获取该tensor的shape
  auto in_dims = engine->getBindingDimensions(engine->getBindingIndex("images")); 
  // auto dd = engine->getTensorShape("images");
  
  iH = in_dims.d[2];
  iW = in_dims.d[3];
  // cout << dd.d[2] << in_dims.d[2] << dd.d[3] << in_dims.d[3] << endl;

  // 计算输入输出tensor需要申请的空间size
  in_size = 1;
  for (int j = 0; j < in_dims.nbDims; j++) {
    in_size *= in_dims.d[j];
  }
  auto out_dims1 = engine->getBindingDimensions(engine->getBindingIndex("num"));
  out_size1 = 1;
  for (int j = 0; j < out_dims1.nbDims; j++) {
    out_size1 *= out_dims1.d[j];
  }
  auto out_dims2 = engine->getBindingDimensions(engine->getBindingIndex("boxes"));
  out_size2 = 1;
  for (int j = 0; j < out_dims2.nbDims; j++) {
    out_size2 *= out_dims2.d[j];
  }
  auto out_dims3 = engine->getBindingDimensions(engine->getBindingIndex("scores"));
  out_size3 = 1;
  for (int j = 0; j < out_dims3.nbDims; j++) {
    out_size3 *= out_dims3.d[j];
  }
  auto out_dims4 = engine->getBindingDimensions(engine->getBindingIndex("classes"));
  out_size4 = 1;
  for (int j = 0; j < out_dims4.nbDims; j++) {
    out_size4 *= out_dims4.d[j];
  }
  context = engine->createExecutionContext(); // 创建一个执行上下文（Execution Context）
  if (!context) {
    cout << "create execution context failed\n";
    std::abort();
  }

  /*
                                  申请空间
  cudaMalloc根据指定的大小，从 GPU 的全局内存中分配一块连续的内存空间，并返回一个指向该内存空间的指针。
  这样就可以在 GPU 上使用该指针来访问和操作所分配的内存空间。
  使用 cudaMalloc() 方法分配的内存需要在使用完毕后通过 "cudaFree()" 方法进行释放，以避免内存泄漏。
  */
  cudaError_t state;
  state = cudaMalloc(&buffs[0], in_size * sizeof(float));
  if (state) {
    cout << "allocate memory failed\n";
    std::abort();
  }
  state = cudaMalloc(&buffs[1], out_size1 * sizeof(int));
  if (state) {
    cout << "allocate memory failed\n";
    std::abort();
  }

  state = cudaMalloc(&buffs[2], out_size2 * sizeof(float));
  if (state) {
    cout << "allocate memory failed\n";
    std::abort();
  }

  state = cudaMalloc(&buffs[3], out_size3 * sizeof(float));
  if (state) {
    cout << "allocate memory failed\n";
    std::abort();
  }

  state = cudaMalloc(&buffs[4], out_size4 * sizeof(int));
  if (state) {
    cout << "allocate memory failed\n";
    std::abort();
  }
  // 通过 cudaStreamCreate() 方法，可以创建一个 CUDA 流对象，并返回一个对应的流句柄（stream handle）。
  // CUDA 流可以用于在 GPU 上并行执行多个操作，例如内存传输、内核函数调用等。通过将不同的操作分配到不同的 CUDA 流中，可以实现这些操作的并发执行，从而提高程序的性能。
  state = cudaStreamCreate(&stream);
  if (state) {
    cout << "create stream failed\n";
    std::abort();
  }
}

letterbox函数用于resize输入图片到模型输入的尺寸，并以灰边填充。这样做的好处是保持原本图片的比例，暴力resize可能会导致图片失真等等。
可以保存中间图片看一下效果：

resize并填充之后
letterbox的思路和做法可以总结为：

确定input和target的size
计算input的宽高和target size的宽高的比值，取最小的那个r。通常是new-shape/shape，因为一般input图片会较大一点，resize成尺寸更小一点的图片。（yolo系列一般都是resize成640）
判断一下图片是否增大（如果原图比resize的尺寸还小，就不用resize了）
先按照比值resize到一个较小的尺寸，（通常有一边到640，有一边小于640）
然后看一下是否要保证resize是32的倍数？
计算短边需要对两边各填充多少
返回1/r

float Yolo::letterbox(
    const cv::Mat& image,  // 输入图片
    cv::Mat& out_image,    // resize填充之后的输出图片
    const cv::Size& new_shape = cv::Size(640, 640),   // resize的size
    int stride = 32,  
    const cv::Scalar& color = cv::Scalar(114, 114, 114),  // 填充的颜色
    bool fixed_shape = false,
    bool scale_up = true) { // scale_up 参数用来控制图片是否增大
  cv::Size shape = image.size();  // 先获取输入图片原本的size
  // cout << "image height: " << shape.height << ", image width: " << shape.width << endl;

  // 计算一下图片原本size和目标size的比例，取小的那个
  float r = std::min(
      (float)new_shape.height / (float)shape.height, (float)new_shape.width / (float)shape.width);

  // scale_up 参数用来控制图片是否增大
  // 如果图片本身比设置的new_size小，需要通过scale_up参数来判断是否resize(放大)，或保持不变
  if (!scale_up) {
    r = std::min(r, 1.0f);
  }

  //cout << "r: " << r << endl;
  // 如果输入图像h:1440，w:2560，resize后会变成h:360，w:640
  int newUnpad[2]{
      (int)std::round((float)shape.width * r), (int)std::round((float)shape.height * r)};


  // 创建一个size为 new shape的 Mat tmp -> resize原图得到的
  cv::Mat tmp;
  if (shape.width != newUnpad[0] || shape.height != newUnpad[1]) {
    cv::resize(image, tmp, cv::Size(newUnpad[0], newUnpad[1]));
  } else {
    tmp = image.clone();
  }

  // 先按最小的比例resize，
  float dw = new_shape.width - newUnpad[0];
  float dh = new_shape.height - newUnpad[1];

  // cout << "newUnpad[0]: " << newUnpad[0] << ", newUnpad[1]: " << newUnpad[1] << endl;
  // cout << "dh1: " << dh << ", dw1: " << dw << endl;

  // fixed_shape 不清楚功能，可能是让填充之后的size能够是32的倍数
  if (!fixed_shape) {
    dw = (float)((int)dw % stride);
    dh = (float)((int)dh % stride); 
    // cout << "dh2: " << dh << ", dw2: " << dw << endl;
  }

  // copyMakeBorder接收参数是向四周扩展的像素数，所以宽高分别除以2
  dw /= 2.0f;
  dh /= 2.0f;
  // cout << "dh3: " << dh << ", dw3: " << dw << endl;


  int top = int(std::round(dh - 0.1f));
  int bottom = int(std::round(dh + 0.1f));
  int left = int(std::round(dw - 0.1f));
  int right = int(std::round(dw + 0.1f));
  cv::copyMakeBorder(tmp, out_image, top, bottom, left, right, cv::BORDER_CONSTANT, color);
  cv::imwrite("letterbox1.jpg", tmp);
  cv::imwrite("letterbox2.jpg", out_image);

  return 1.0f / r;
}

blobFromImage函数就是将img拉长放在一个float数组里，且标准化。draw_objects函数将检测结果绘制在图片上，用到的也都是一些opencv的函数。

Infer函数是模型执行推理的函数，主要就是遵循TRT和CUDA那一套流程：

用letterbox对输入图片进行不失真的resize
图片从BGR转RGB后，做blob返回一个数组（opencv读取图片好像默认是BGR的）
为输出的结果申请buffer
把input转换成的数组传输到GPU上
执行推理
把output从GPU传输到Host上
把output转换到对应原图size结果，并把结果写出传入参数里
释放blob空间

void Yolo::Infer(
    int aWidth,     //图片的宽
    int aHeight,    //图片的高
    int aChannel,   //图片的通道数
    unsigned char* aBytes,    //图片的数据 cv::Mat data值
    float* Boxes,    // 用于存放结果的数组指针 
    int* ClassIndexs,
    int* BboxNum) {
  cv::Mat img(aHeight, aWidth, CV_MAKETYPE(CV_8U, aChannel), aBytes); // 输入的图片
  cv::Mat pr_img; // 转换后的图片先预定义
  float scale = letterbox(img, pr_img, {iW, iH}, 32, {114, 114, 114}, true); // 通过letterbox进行转换
  cv::cvtColor(pr_img, pr_img, cv::COLOR_BGR2RGB); // 图片BGR转RGB
  float* blob = blobFromImage(pr_img); //执行blob

  // 静态变量
  static int* num_dets = new int[out_size1]; // 对应检测到的每个目标
  static float* det_boxes = new float[out_size2]; // 每个目标对应的box
  static float* det_scores = new float[out_size3]; // 每个目标的置信度得分 
  static int* det_classes = new int[out_size4]; // 每个目标的类别

  // cudaMemcpyAsync 是 CUDA 库中的一个函数，用于在 GPU 之间或在主机和设备之间异步地进行内存数据的传输。
  // 使用 cudaMemcpyAsync 函数进行内存数据传输时，数据传输是异步执行的，即函数调用会立即返回，而不会等待数据传输完成。
  // 因此，可以在数据传输进行的同时进行其他的 GPU 计算操作，以提高程序的效率。
  // 这里相当于把blob上的数据复制到buffs上
  cudaError_t state =
      cudaMemcpyAsync(buffs[0], &blob[0], in_size * sizeof(float), cudaMemcpyHostToDevice, stream);
  if (state) {
    cout << "transmit to device failed\n";
    std::abort();
  }

  // enqueueV2 是 nvinfer1::IExecutionContext::enqueueV2方法，有三个参数
  // bindings 输入或输出array的缓存的指针
  // stream  一个cuda stream
  // inputConsumed	 输入的buffer可以被重新填充的信号(不清楚怎么使用)
  // 个人理解就是执行推理
  context->enqueueV2(&buffs[0], stream, nullptr);

  // 推理结束之后，把device上的output数据再传回host
  state =
      cudaMemcpyAsync(num_dets, buffs[1], out_size1 * sizeof(int), cudaMemcpyDeviceToHost, stream);
  if (state) {
    cout << "transmit to host failed \n";
    std::abort();
  }
  state = cudaMemcpyAsync(
      det_boxes, buffs[2], out_size2 * sizeof(float), cudaMemcpyDeviceToHost, stream);
  if (state) {
    cout << "transmit to host failed \n";
    std::abort();
  }
  state = cudaMemcpyAsync(
      det_scores, buffs[3], out_size3 * sizeof(float), cudaMemcpyDeviceToHost, stream);
  if (state) {
    cout << "transmit to host failed \n";
    std::abort();
  }
  state = cudaMemcpyAsync(
      det_classes, buffs[4], out_size4 * sizeof(int), cudaMemcpyDeviceToHost, stream);
  if (state) {
    cout << "transmit to host failed \n";
    std::abort();
  }
  BboxNum[0] = num_dets[0];
  int img_w = img.cols;
  int img_h = img.rows;
  // cout << "iW: " <<iW << "iH: " <<iH << "img_w: " <<img_w << "img_h: " <<img_h << "scale: " <<scale << endl;

  // 计算一下宽高不同的图片转换回来的偏差，因为上面的letterbox是按最小的scale缩放的
  int x_offset = (iW * scale - img_w) / 2;
  int y_offset = (iH * scale - img_h) / 2;
  for (size_t i = 0; i < num_dets[0]; i++) {
    float x0 = (det_boxes[i * 4]) * scale - x_offset;
    float y0 = (det_boxes[i * 4 + 1]) * scale - y_offset;
    float x1 = (det_boxes[i * 4 + 2]) * scale - x_offset;
    float y1 = (det_boxes[i * 4 + 3]) * scale - y_offset;

    // 防止框的点越界
    x0 = std::max(std::min(x0, (float)(img_w - 1)), 0.f);
    y0 = std::max(std::min(y0, (float)(img_h - 1)), 0.f);
    x1 = std::max(std::min(x1, (float)(img_w - 1)), 0.f);
    y1 = std::max(std::min(y1, (float)(img_h - 1)), 0.f);
    Boxes[i * 4] = x0;
    Boxes[i * 4 + 1] = y0;
    Boxes[i * 4 + 2] = x1 - x0;
    Boxes[i * 4 + 3] = y1 - y0;
    ClassIndexs[i] = det_classes[i];
  }
  delete blob;
}

析构函数就是一些内存、实例的销毁

Yolo::~Yolo() {
  cudaStreamSynchronize(stream); // 用于同步等待指定的 CUDA 流上的所有操作完成。
  cudaFree(buffs[0]);
  cudaFree(buffs[1]);
  cudaFree(buffs[2]);
  cudaFree(buffs[3]);
  cudaFree(buffs[4]);
  cudaStreamDestroy(stream);
  context->destroy();
  engine->destroy();
  runtime->destroy();
}

最后就是主程序调用实验了，有两个入参：trt模型路径和需要推理的图片路径

int main(int argc, char** argv) {
  // 输入参数有-model_path 和 -image_path
  if (argc == 5 && std::string(argv[1]) == "-model_path" && std::string(argv[3]) == "-image_path") {
    char* model_path = argv[2];
    char* image_path = argv[4];
    float* Boxes = new float[4000];
    int* BboxNum = new int[1];
    int* ClassIndexs = new int[1000];
    Yolo yolo(model_path);
    cv::Mat img;
    img = cv::imread(image_path);
    // warmup 
    for (int num =0; num < 10; num++) {
      yolo.Infer(img.cols, img.rows, img.channels(), img.data, Boxes, ClassIndexs, BboxNum);
    }
    // run inference
    auto start = std::chrono::system_clock::now();
    yolo.Infer(img.cols, img.rows, img.channels(), img.data, Boxes, ClassIndexs, BboxNum);
    auto end = std::chrono::system_clock::now();
    std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;

    yolo.draw_objects(img, Boxes, ClassIndexs, BboxNum);

  } else {
    std::cerr << "--> arguments not right!" << std::endl;
    std::cerr << "--> yolo -model_path ./output.trt -image_path ./demo.jpg" << std::endl;
    return -1;
  }
}

结语

回忆了很多c++的基本知识，也对齐了python用TRT推理和c++用TRT的一个流程上的异同，并且在查找很多TRT的c++ api时发现很多api弃用了或者推荐使用新的api，后续计划将该工程对应的api修改换成新版的作为一个课后补习。

如有错误，请联系我修改。

sansheng0208

关注

18
点赞
踩
27

收藏

觉得还不错? 一键收藏
1
评论
Yolo系列模型使用TensorRT-C++推理学习

回忆了很多c++的基本知识，也对齐了python用TRT推理和c++用TRT的一个流程上的异同，并且在查找很多TRT的c++ api时发现很多api弃用了或者推荐使用新的api，后续计划将该工程对应的api修改换成新版的作为一个课后补习。如有错误，请联系我修改。
复制链接

扫一扫