1、前言
在前两篇博客中学习了yolo&tensorrt推理的代码，学习了模型的加载、预处理、后处理等等，本文将继续学习其中的消费者、生产者推理代码
yolo&tensorrt项目:https://github.com/shouxieai/infer
第一部分学习记录:https://blog.csdn.net/weixin_42108183/article/details/129411759
第二部分学习记录:https://blog.csdn.net/weixin_42108183/article/details/129455120
2、多线程代码学习
cpm.hpp
// 结果类型、输入类型、模型类型
template <typename Result, typename Input, typename Model>
class Instance{
protected:
// 任务形式
struct Item
{
Input input; // 输入
std::shared_ptr<std::promise<Result>> pro; // 输入对应的输出
};
std::condition_variable cond_; // 条件变量
std::queue<Item> input_queue_; // 任务队列
std::shared_ptr<std::thread> worker_; // 工作线程
volatile bool run_ = false; // 是否运行
public:
// 析构函数中使用stop回收子线程，即推理线程
virtual ~Instance() {
stop(); }
void stop(){
run_ = false; // 将运行状态设置为false
cond_.notify_one(); // 唤醒工作线程、如果工作线程检测到run_=false,则会退出。
{
std::unique_lock<std::mutex> l(queue_lock_);
// 如果工作队列非空,继续推理完毕
while (!input_queue_.empty()){
auto &item = input_queue_.front();
if (item.pro)
item.pro->set_value(Result());
input_queue_.pop();
}
}
// 回收工作线程
if (worker_)
{
worker_->join();
worker_.reset();
}
}
// 推理单张图像
virtual std::shared_future<Result> commit(const Input &input){
// 构建任务
Item item;
item.input = input;
item.pro.reset(new std::promise<Result>());
{
// 给任务队列添加任务
std::unique_lock<std::mutex> __lock_(queue_lock_);
input_queue_.push(item);
}
// 唤醒线程
cond_.notify_one();
return item.pro->get_future();
}
// 推理多张图片
commits(const std::vector<Input> &inputs){
}
template <typename LoadMethod>
bool start(...){
// 先stop，保证没有子线程
stop();
std::promise<bool> status; // 模型是否被正确加载的标识
// 创建工作线程
worker_ = std::make_shared<std::thread>(&Instance::worker<LoadMethod>, this,
std::ref(loadmethod), std::ref(status));
return status.get_future().get(); // 阻塞等待 status被设置某个值
}
private:
template <typename LoadMethod>
void worker(const LoadMethod &loadmethod, std::promise<bool> &status){
// 加载模型
std::shared_ptr<Model> model = loadmethod();
if (model == nullptr)
{
// 加载失败返回false
status.set_value(false);
return;
}
// 等待解锁
while (get_items_and_wait(fetch_items, max_items_processed_)){
inputs.resize(fetch_items.size());
std::transform(fetch_items.begin(), fetch_items.end(), inputs.begin(),
[](Item &item)
{
return item.input; });
// 推理
auto ret = model->forwards(inputs, stream_);
}
model.reset(); // 智能指针，释放模型
run_ = false;
}
virtual bool get_items_and_wait(std::vector<Item> &fetch_items,int max_size){
// 是否解锁
std::unique_lock<std::mutex> l(queue_lock_);
// run_=false 或者 任务队列不为空，则解锁
cond_.wait(l, [&](){
return !run_ || !input_queue_.empty(); });
for (int i = 0; i < max_size && !input_queue_.empty(); ++i)
{
// 添加任务
fetch_items.emplace_back(std::move(input_queue_.front()));
input_queue_.pop();
}
return true;
}
// 推理单张图片时使用
virtual bool get_item_and_wait(Item &fetch_item){
}
}
3、测试代码
main.cpp
void perf()
{
// 将图片添加到batchsize
for (int i = images.size(); i < batch; ++i)
images.push_back(images[i % 3]);
// 初始化消费者、生产者实例
cpm::Instance<yolo::BoxArray, yolo::Image, yolo::Infer> cpmi;
// start中会开辟子线程用于加载engine,此时任务队列空，子线程会阻塞
bool ok = cpmi.start(
[]
{
return yolo::load("yolov8n.transd.engine", yolo::Type::V8); },
max_infer_batch);
// 将images进行cvimg 的操作，然后储存到yoloimages中
td::transform(images.begin(), images.end(), yoloimages.begin(), cvimg);
for (int i = 0; i < 5; ++i)
{
timer.start();
// 提交图片，.get()会阻塞等待子线程返回结果
cpmi.commits(yoloimages).back().get();
timer.stop("BATCH16");
}
...
}
4、总结
本次学习了基于生产者、消费者的多线程推理代码，并且很容易的将代码融入到项目中。