Yolov5实例分割Tensorrt部署实战

Posted 2022-12-04 机器鱼

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 [e-mail]:ultralytics/yolov5.git#官方代码
git clone [e-mail]: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点个心，谢谢