在这个项目中,我们实现了从 YOLO 的 onnx 模型到 tensorrt 模型的转变,并在其中加入了 NMS tenosrrt 的插件编写,使得网络的速度提升了很多。
pip install --upgrade setuptools pip --user
pip install nvidia-pyindex
pip install --upgrade nvidia-tensorrt
pip install pycuda
1、下载 yolov7, 并导出 onnx。
1.1、下载源码
git clone https://github.com/WongKinYiu/yolov7.git
1.2、下载权重
权重再 yolov7 的官方仓库里面有链接,大家可以选择自己的想要转化的权重, 都是OK 的。
wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7x.pt
1.3、安装所需要的包
pip install -r yolov7/requirements.txt
上面这一步,一般情况下,不要直接安装,有可能会破坏你原有的环境,点开 requirements.txt ,看自己缺少什么文件包,再依次安装就好了。当然要是你喜欢直接一点的,那一下子全部安装也是可以的。
1.4、导出 onnx 模型
python yolov7/export.py --weights yolov7-tiny.pt --grid --simplify
--grid --simplify
可以将
2、运行 export_TRT.py
目前里面是包含了写了代码,加载了 TensorRT 官方的 NMS 的 plugin,所以支持 end2end 的操作。
2.1、不包含 NMS plugin
python export_TRT.py -o weights/yolov7x.onnx
2.2、包含 NMS plugin
python export_TRT.py -o weights/yolov7x.onnx -s --end2end
也可直接运行:examples.ipynb 文件
说明 :
1、export_TRT.py 可以在运行终端输入很多参数,但大多数的参数都包含了默认值,想要修改的话,直接在终端输入对应的参数的即可。
2、默认运行 导出 TensorRT 模型,并测试网络性能,到处的 TensorRT 模型和输出的 ONNX 模型位置相同,名字相同,仅仅是后缀名不相同。
3、代码中包含 fps 测试模块,可以对于使用 NMS plugin 和不使用 NMS plugin 的 fps 上的差别;
在 images 文件下,有一个时评 demo,检测代码的时候,可以尝试运行这个视频
参考大神仓库:https://github.com/Linaom1214/TensorRT-For-YOLO-Series