Jetson Xavier NX 配置(五)—— Yolo5+TensorRT加速的pytorch环境配置

1、Yolov5跑个demo

        下载yolov5的工程文件

git clone -b v5.0 https://github.com/ultralytics/yolov5.git 

        下载.pt模型文件,一定要注意版本一致!Releases · ultralytics/yolov5 · GitHub ,此处下载最小的模型yolov5s.pt

        用demo中的图片测试一下,单帧82ms:

python detect.py --source bus.jpg --weights yolov5s.pt --conf 0.25

Jetson Xavier NX 配置(五)—— Yolo5+TensorRT加速的pytorch环境配置_第1张图片

        以上过程中如果报错缺哪些库,缺啥装啥

2、TensorRT加速

简单思路:根据.pt模型文件生成.wts文件,再生成.engine文件

需要的库

        pycuda

        tensorrt

注:由于Jetpack组件装完以后自动在本地安装过tensorrt,因此在python虚拟环境中运行的话,只需将本地的tensorrt两个文件夹拷贝到虚拟环境的文件夹中即可(前提是俩环境的python版本一致)。(以下仅是我的存放路径,仅供参考)

  • 从本地拷贝出来的路径:/usr/lib/python3.6/dist-packages
  • 要拷贝到虚拟环境的路径:/home/yishengjun/miniforge3/envs/forTorch/lib/python3.6/site-packages  (其中yishengjun是用户名,forTorch是python虚拟环境的名)

需要的工程:(此处选择yolov5 v5.0版本的)

git clone -b v5.0 https://github.com/ultralytics/yolov5.git 
git clone -b yolov5-v5.0 https://github.com/wang-xinyu/tensorrtx.git

        注意!!!很容易下错版本,一定注意yolov5和tensorrtx的版本要一致!要么都v4.0要么都v5.0!否则最后生成.engine文件的时候会报错什么什么"Assertion 'scale_1' failed",大概率是版本不一致的问题。

步骤

        把yolov5s.pt放到yolov5工程的weights文件夹下
        把tensorrtx工程yolov5文件夹下的gen_wts.py复制到yolov5工程下
        进入yolov5工程,执行代码,实现:根据.pt文件生成.wts文件

python3 gen_wts.py -w weights/yolov5s.pt

        在tensorrtx工程的yolov5文件夹下,新建build文件夹
        在build文件夹下进入terminal,执行代码,实现:根据.wts文件生成.engine文件

sudo ./yolov5 -s yolov5s.wts yolov5s.engine s

        ⑥ 运行TensorRT加速后的yolov5

sudo ./yolov5 -d yolov5s.engine ../samples/bus.jpg

Jetson Xavier NX 配置(五)—— Yolo5+TensorRT加速的pytorch环境配置_第2张图片

之前图像检测一帧约60ms~80ms左右,TensorRT加速后30ms+,可以做实时的应用了

参考:Jetson AGX Xavier实现TensorRT加速YOLOv5进行实时检测_围白的尾巴-CSDN博客

你可能感兴趣的:(Jetson,Xavier,NX,python,开发语言,后端)