8.5.tensorRT高级(3)封装系列-基于生产者消费者实现的yolov5封装
目录
-
- 前言
- 1. yolov5封装
- 总结
前言
杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。
本次课程学习 tensorRT 高级-基于生产者消费者实现的yolov5封装
课程大纲可看下面的思维导图
1. yolov5封装
这节我们学习使用封装好的组件，配合生产者消费者，实现一个完整的 yolov5 推理
对于 yolov5 的封装，主要考虑以下：
1. 希望调用者是线程安全的，可以随意进行 commit，而不用考虑是否冲突
2. 希望结果是懒加载的，也就是需要的时候才等待，不需要的时候可以不等待
- 由 promise 与 future 配合实现
- 这样的灵活度和效率性能都是最好的
3. 希望最大化利用 GPU，如何利用呢？需要尽可能的使得计算密集
- 实际体现就是抓取一个批次，一次性进行推理
- 这在内部的消费者模型里面，一次抓一批
我们来看代码：
yolov5.hpp
#ifndef YOLOV5_HPP
#define YOLOV5_HPP
#include yolov5.cpp
#include "yolov5.hpp"
#include 头文件中定义了一个 Infer 类，该类只有一个 commit 纯虚函数，接收图像数据，然后 shared_future 对象，create_infer 函数用于 RAII，它创建并返回一个 Infer 接口类的实现，通过 startup 方法初始化实例
在 startup 函数中我们创建了一个 bool 类型的 promise 变量，用于判断资源是否获取成功，通过引用的方式传递到了消费者线程 worker 中，在 worker 线程里面会处理模型加载的过程，加载成功或者失败都会把结果反馈到对应的 promise 变量上，而通过 future 对象的 get 方法我们可以获取到是否成功，而成功之后消费者线程也启动了，会继续往下走
在 worker 消费者线程中，有个条件变量在等待，如果条件为 true 则退出等待，也就是当队列不为空时会退出等待，然后将队列中的数据移动到 vector 容器中，循环 vector 容器中的每个图像进行 brg2rgb 同时 split，然后将这一批数据送到网络进行推理拿到结果，然后对拿到的 box 进行decode，最好将通过 promise 的 set_value 方法将 boxes 返回回去
我们再来看下 commit 生产者线程，它会将接收的图像数据进行预处理后放入到队列中，然后利用条件变量通知消费者线程可以处理了，其中预处理使用的是 warpAffine 仿射变换，
总的来说，上述代码提供了 YOLOv5 的推理功能。它实现了一个生产者-消费者模式，其中生产者可以异步地提交推理任务，并使用 future-promise 获取结果。worker 线程作为消费者执行这些任务。以下是 yolov5 封装的关键点：
1. 封装的接口：定义了 Infer 接口，提供异步推理的方法
2. 生产者-消费者模式：该实现采用生产者-消费者模式，其中生产者提交推理任务，worker 线程作为消费者去执行
3. 异步处理：使用 future-promise 机制，生产者可以异步地提交任务并等待结果
4. 预处理：预处理采用 warpAffine 仿射变换
5. decode：采用 IM 逆矩阵进行解码恢复成框
6. 多线程安全：使用互斥锁和条件变量确保多线程安全
7. 资源管理：使用 RAII 原则确保资源的正确管理
最好我们来看下 main.cpp 中的内容变化：
// tensorRT include
// 编译用的头文件
#include 模型构建部分非常简单，一行代码解决模型编译问题，非常方便。推理部分通过 create_infer 拿到一个 shared_ptr 对象，通过调用 commit 方法把图像加入到生成者队列中，然后通过 get 方法等待消费者线程拿到推理结果，然后直接绘制目标框即可，相比于之前的纯裸的，无封装的 yolov5 来说非常简便了
总结
本次课程学习了基于生产者和消费者实现的 yolov5 封装，commit 生成者线程不断往队列中抛数据，通过条件变量 cv_ 通知消费者进行消费，worker 消费者线程会一直等待，直到队列不为空，它会把一批次数据全部扔到模型中进行推理，然后解码，通过 promise 将结果传递回来。create_infer 是 RAII 和接口模式的体现，通过 startup 获取资源并初始话，如果资源获取失败则直接退出，同时实例化的是实现类 InferImpl，而返回的是接口类 Infer，只对使用者暴露 commit 接口。这都是我们之前在生产者消费者课程中讲到过的知识，这边是直接拿过来用了。