在阅读本章内容之前,需要把部署环境以及训练环境都安装好。
1.TensorRTX下载
这里使用Wang-xinyu大佬维护的TensorRTX库来对YOLOv8进行推理加速的演示,顺便也验证一下前面环境配置的成果。
github地址:GitHub - wang-xinyu/tensorrtx,下载后放到wsl的路径下,我这里放在/home路径下,并解压到/home路径下:
unzip tensorrtx-master.zip
2.Vs-Code远程连接TensorRTX
首先必须先安装一个WSL插件,如下所示:
然后连接WSL-Ubuntu2004,如下所示:
最后打开我们解压的tensorRTX文件,因为我们需要用到YOLOv8,所以选择YOLOv8,然后点击OK就可以了。
因为后面同时涉及到Python代码和C++代码,所以把这两个用到的插件也一起安装了,其中python环境需要安装Python插件,C++环境需要安装c++插件.
3.使用YOLOv8n.pt进行推理加速实验
这里就先使用官方的YOLOv8n.pt目标检测模型进行实验,看看我们的环境是否安装成功,这里只是先走一遍流程,关于TensorRTX的具体细节后面会有专门的文章来介绍。
第一步:当然是要把官方的YOLOv8n.pt放到/home/tensorrtx-master/yolov8路径下,然后执行gen_wts.py代码,修改pt_file路径(-w)和生成的wts_file路径(-o)就行,如下所示。这段代码是把.pt模型的参数大小及权值保存下成.wts,后面也是根据该.wts文件的保存规则解析并序列化成.engine文件。
python gen_wts.py -w /home/tensorrtx-master/yolov8/yolov8n.pt -o /home/tensorrtx-master/yolov8/yolov8n.wts -t detect
第二步:需要配置相关的头文件,shift+ctrl+P,搜索Edit configurations (UI),然后在Include path这里补充上以下路径,也就是我们前面安装的Cuda、OpenCV和TensorRT的头文件路径。
/usr/local/cuda-11.6/include
/usr/local/include/opencv4
/root/TensorRT-8.5.1.7/include
同时需要修改一下CMakeLists.txt文件中关于tensorRT的include和lib路径,如下所示.
然后就可以在vs-code终端执行以下指令:
mkdir build
cd build
cp /home/tensorrtx-master/yolov8/yolov8n.wts /home/tensorrtx-master/yolov8/build
cmake ..
make
执行完之后发现在build路径下会生成可执行文件yolov8_det.
第三步:序列化生成.engine模型,这一步时间会稍微长一点,执行成功后会在build路径下生成yolov8n.engine文件。
./yolov8_det -s /home/tensorrtx-master/yolov8/yolov8n.wts yolov8n.engine n
第四步:反序列化.engine模型并进行推理,推理结果也是直接保存在build路径下,
./yolov8_det -d yolov8n.engine ../images c
我用多张图片进行测试,推理时间是很快的,只花费了1ms,如下所示。
至此,使用WSL-Ubuntu20.04环境对YOLOv8模型进行了推理加速的实验,但是我们发现WSL-Ubuntu20.04配置环境也还是很复杂,要手动安装那么多东西,下面我会介绍环境搭建神器Docker,来快速搭建部署及训练环境,实现Develop faster. Run anywhere.这种效果。