文章目录

  • 前言
  • 1.Yolo简介
  • 2.onnxruntime简介
  • 3.Yolov5模型训练及转换
  • 4.利用cmake向C++部署该onnx模型
  • 总结

前言

接到一个项目,需要用c++和单片机通信,还要使用yolo模型来做到目标检测的任务,但目前网上的各种博客并没有完整的流程教程,让我在部署过程费了不少劲,也踩了不少坑(甚至一度把ubuntu干黑屏)。于是想把训练及部署过程记录下来,并留给后来者方便使用。(博主使用的系统是ubuntu20.04)

1.Yolo简介

作为一个经典且实用的目标检测模型,yolo的性能强大已无需多言,现在(2023.4.1)yolo模型已经推出到yolov8,但是推理速度上yolov5还是够快且够用,而且对各种外部硬件的适配性强,如oak相机支持的模型还在yolov5-yolov6,所以博主这里选择yolov5进行训练。

2.onnxruntime简介

onnxruntime是微软推出的一款推理框架,我们可以很方便的利用它运行一个onnx模型,而且它支持多种运行后端,包括CPU,GPU,TensorRT,DML等。onnxruntime可以说是对onnx模型最原生的支持了,而且onnxruntime也有在C++上部署使用的相关库,所以我们选择onnxruntime作为我们的推理框架进行部署。

3.Yolov5模型训练及转换

怎么训练呢?

对于yolov5模型的训练,其实选择哪个都可以,博主这里模型使用的是https://github.com/ultralytics/yolov5/tree/v5.0

  • 训练过程参考这个链接
    这个博主对训练过程提供了非常详细的说明。按照这个博主的教程可以得到训练后的.pt文件,先别急着在models里export转换模型。如果在这里转换之后得到的output可能是这样的。

可以看到这个图片中有三个output,这样的结果是不能用的,我们需要的output是类似于[1,25200,7]这样的结果,并且这个结果必须在第一位上,通过观察可以知道25200是上图中三个输出结果中三个先验框大小和数量相乘得到的而我们训练使用的yolov5模型export的转换结果却不是这样的总和。

那怎么转换呢?

很简单,通过观察yolov5中yolo部分的代码发现


这个代码中的return是有问题的,返回结果不是我们想要的存在类似与[1,25200,7]的结果,而是详细细分的。这里我们最好不要改动yolo部分的代码,这可能会牵连其他部分从而产生错误。我们直接使用https://github.com/ultralytics/yolov5 这个模型来进行转换

  • 首先先下载这个模型,用vscode打开后找到export.py
    在vscode的终端中输入
python export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...

注意yolov5s.pt处填自己训练结果的pt文件地址。然后我们使用netron查看得到的onnx模型

  • 惊奇的发现我们得到了我们想要的output!(win!)(注意最后的85只是80个labels和5个参数,自己训练的模型有自己设置的labels所以数量可能不一样)

  • 到这里我们已经完成了yolov5模型的训练和转换,请带好自己的随身物品(onnx模型),我们要赶往下一站了

4.利用cmake向C++部署该onnx模型

  • 在这一步我们可能会遇到很多难题,但是不要着急,我来带着你一步步走

1.下载onnxruntime-gpu版本

  • 博主这里使用的是onnxruntime-gpu-1.7

  • 这里选择onnxruntime-linux-x64-gpu-1.7.0.tgz
  • 下载后解压,然后放在主目录里,我们可以看到打开文件后有include和lib两个文件夹。这个我们之后要在cmakelist里链接。

2.下载vscode并创建cmake工程

  • 这是博主的main函数代码,里面包含了onnxruntime的使用和推理结果的处理
#pragma comment(lib, "k4a.lib")#include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include "opencv2/imgproc/imgproc_c.h"#include//onnxruntime// #include // #include // #include // #include #include #include #include // 命名空间using namespace std;using namespace cv;using namespace Ort;// 自定义配置结构struct Configuration{public: float confThreshold; // Confidence threshold置信度阈值float nmsThreshold;// Non-maximum suppression threshold非最大抑制阈值float objThreshold;//Object Confidence threshold对象置信度阈值string modelpath;};// 定义BoxInfo结构类型typedef struct BoxInfo{float x1;float y1;float x2;float y2;float score;int label;} BoxInfo;// int endsWith(string s, string sub) {// return s.rfind(sub) == (s.length() - sub.length()) " /> input_image_;// 输入图片void normalize_(Mat img);// 归一化函数void nms(vector& input_boxes);Mat resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left);Env env = Env(ORT_LOGGING_LEVEL_ERROR, "yolov5-6.1"); // 初始化环境Session *ort_session = nullptr;// 初始化Session指针选项SessionOptions sessionOptions = SessionOptions();//初始化Session对象//SessionOptions sessionOptions;vector input_names;// 定义一个字符指针vectorvector output_names; // 定义一个字符指针vectorvector<vector> input_node_dims; // >=1 outputs,二维vectorvector<vector> output_node_dims; // >=1 outputs ,int64_t C/C++标准};YOLOv5::YOLOv5(Configuration config){this->confThreshold = config.confThreshold;this->nmsThreshold = config.nmsThreshold;this->objThreshold = config.objThreshold;this->num_classes = 1;//sizeof(this->classes)/sizeof(this->classes[0]);// 类别数量this->inpHeight = 640;this->inpWidth = 640;string model_path = config.modelpath;//std::wstring widestr = std::wstring(model_path.begin(), model_path.end());//用于UTF-16编码的字符//gpu, https://blog.csdn.net/weixin_44684139/article/details/123504222//CUDA加速开启//OrtSessionOptionsAppendExecutionProvider_Tensorrt(sessionOptions, 0);OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0);sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);//设置图优化类型//ort_session = new Session(env, widestr.c_str(), sessionOptions);// 创建会话,把模型加载到内存中//ort_session = new Session(env, (const ORTCHAR_T*)model_path.c_str(), sessionOptions); // 创建会话,把模型加载到内存中ort_session = new Session(env, (const char*)model_path.c_str(), sessionOptions);size_t numInputNodes = ort_session->GetInputCount();//输入输出节点数量 size_t numOutputNodes = ort_session->GetOutputCount(); AllocatorWithDefaultOptions allocator; // 配置输入输出节点内存for (int i = 0; i GetInputName(i, allocator));// 内存Ort::TypeInfo input_type_info = ort_session->GetInputTypeInfo(i); // 类型auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();// auto input_dims = input_tensor_info.GetShape();// 输入shapeinput_node_dims.push_back(input_dims);// 保存}for (int i = 0; i GetOutputName(i, allocator));Ort::TypeInfo output_type_info = ort_session->GetOutputTypeInfo(i);auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();auto output_dims = output_tensor_info.GetShape();output_node_dims.push_back(output_dims);}this->inpHeight = input_node_dims[0][2];this->inpWidth = input_node_dims[0][3];this->nout = output_node_dims[0][2];// 5+classesthis->num_proposal = output_node_dims[0][1];// pre_box}Mat YOLOv5::resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left)//修改图片大小并填充边界防止失真{int srch = srcimg.rows, srcw = srcimg.cols;*newh = this->inpHeight;*neww = this->inpWidth;Mat dstimg;if (this->keep_ratio && srch != srcw) {float hw_scale = (float)srch / srcw;if (hw_scale > 1) {*newh = this->inpHeight;*neww = int(this->inpWidth / hw_scale);resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);*left = int((this->inpWidth - *neww) * 0.5);copyMakeBorder(dstimg, dstimg, 0, 0, *left, this->inpWidth - *neww - *left, BORDER_CONSTANT, 114);}else {*newh = (int)this->inpHeight * hw_scale;*neww = this->inpWidth;resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);//等比例缩小,防止失真*top = (int)(this->inpHeight - *newh) * 0.5;//上部缺失部分copyMakeBorder(dstimg, dstimg, *top, this->inpHeight - *newh - *top, 0, 0, BORDER_CONSTANT, 114); //上部填补top大小,下部填补剩余部分,左右不填补}}else {resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);}return dstimg;}void YOLOv5::normalize_(Mat img)//归一化{//img.convertTo(img, CV_32F);//cout<<"picture size"<

  • 上图是博主的驱动。

2.cuda和cudnn的安装

  • 博主这里选择的版本是cuda11.2版本
  • 参考https://blog.csdn.net/weixin_42156097/article/details/127321349 这个博文详细记录了CUDA和cudnn的安装过程,记得版本不要搞错,cudnn的版本也要对应cuda和nvidia来整。
  • 如果你安装完了cuda和cudnn,就可以开始快乐的cmake了!

3.cmake

  • 我们就不详细介绍cmake相关的内容了,按照上面的代码放进去后就可以直接cmake . make ./yourproject了。
  • 然后就是快乐的结果展示时间辣!

  • 可以看到每一帧的处理速度在0.06s左右,可以达到实时监测辣!
  • 后续需要添加什么功能就可以自己在c++中改辣!不喜欢用python处理数据的人的福音

总结

  • 博主为了个通信的功能以及不想用python写代码开启了这段部署之旅,中间经历了onnxruntime安装后不会cmake链接,链接后发现缺东少西,甚至连nvidia的驱动都没装就开始调用cuda的gpu加速,装nvidia驱动甚至差点把ubuntu搞坏。这段血与泪我不想让第二个人经历了,所以写下这个博客来记录和分享。

之后博主遇到的问题 (2023.4.8)

  • 在cuda和cudnn版本选择上博主后来用的cuda版本是11.0.3,onnxruntime版本换成了18。原来那个版本在写下第一版博文后过了几天就出现了版本问题。
  • nvidia驱动安装后最好执行sudo apt-mark hold nvidia-driver-525版本改成你自己的。这个命令可以阻止驱动自我更新,不然如果更新后会出现两个版本的驱动,出现系统版本和自己装的驱动版本不同从而导致错误的情况。

  • 之后我可能会做一些其他方面的博客,如果有需要的话就尽情留言吧!

我是一只咸鱼_,一个做视觉的小白,文章对你有帮助的话就点赞收藏吧! kisskiss!