0.引言
ultralytics在yolov5的6.2版本发布了实例分割模型,可实现快速实例分割,采用官方v5s-seg.pt效果如下图所示:
本博客将以此为基础,开发c++版本的tensorrt推理代码,直接上链接:这里,本人环境为:
cuda10.2 cudnn8.2.4 Tensorrt8.0.1.6 Opencv4.5.4。。。。。。代码列表如下
├── CMakeLists.txt
├── images
│ ├── bus.jpg
│ └── zidane.jpg
├── logging.h
├── main1_onnx2trt.cpp
├── main2_trt_infer.cpp
├── models
│ ├── yolov5s-seg.engine
│ └── yolov5s-seg.onnx
├── output.jpg
├── README.md
└── utils.h
1.生成onnx模型
首先我们克隆最新的版本的代码,即6.2版本,并下载对应的pt模型,这里以yolov5s-seg.pt为例进行后面的叙述。
git clone git@github.com:ultralytics/yolov5.git#官方代码
git clone git@github.com:fish-kong/Yolov5-instance-seg-tensorrt.git#我的tensort推理c++代码
yolov5-6.2的官方代码export.py中提供了直接生成engine的方法,但是我不推荐直接用,原因是生成的engine是跟电脑环境有关的,你换了一个环境之后,之前电脑生成的engine就不能用了,除非两个电脑的环境一模一样,所以我们仅生成onnx模型,命令如下
python export.py --data coco128-seg.yaml --weights yolov5s-seg.pt --cfg yolov5s-seg.yaml --include onnx
将会生成名字叫yolov5s-seg.onnx,采用Netron打开后,我们可以看到输入是1x3x640x640,输出output0是1x25200x117,输出output1是1x32x160x160,这些size对后续推理很重要,都是需要写进c++推理代码中的参数。
2. 生成engine模型
1.首先定位到你clone的repo目录下,就是Yolov5-instance-seg-tensorrt目录下
2.复制 yolov5s-seg.onnx 到 models/
3.运行下列代码,生成转换与推理的可执行文件-->onnx2trt trt_infer
mkdir build
cd build
cmake ..
make
4.模型转换
sudo ./onnx2trt ../models/yolov5s-seg.onnx ../models/yolov5s-seg.engine
通过上述操作,我们能得到yolov5s-seg.engine(前提是安装好了cuda,cudnn,tensorrt,opencv,推荐采用我的版本)
3.推理
通过2节操作,其实已经生成了trt_infer这个可执行文件只需要执行即可
sudo ./trt_infer ../models/yolov5s-seg.onnx ../images/bus.jpg
for (int i = 0; i < 10; i++) {//计算10次的推理速度
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, prob1, 1);
auto end = std::chrono::system_clock::now();
std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
}
在main2_trt_infer.cpp中的上述代码是推理10次,并显示时间,在我的1080ti上,基本维持在10ms一帧图片,速度相当之快。
最终结果如下,对比文章最前面的图(这个是官方code直接用pt推理的结果),可以看出基本一致。
4.参考
1.wangxinyu大佬的tensorrtx
2.UNeedCryDear大佬的opencv推理
3.2022.09.29更新 c++下面使用opencv部署yolov5和yolov7实例分割模型(六)_爱晚乏客游的博客-CSDN博客_c++ opencv yolo分割
完整的代码已经上传,直接clone使用即可,欢迎评论区评论如果觉得有用,麻烦给我的github repo点个心,谢谢