部署YOLOv5算法在NVIDIA Jetson Nano上并使用TensorRT和DeepStream进行加速涉及几个关键步骤。下面是一个详细的指南:
步骤 1: 准备YOLOv5模型
训练或下载预训练模型:首先,你需要有一个YOLOv5模型。你可以自己训练一个模型,或者使用从YOLOv5的GitHub仓库下载的预训练模型。
git clone -b v7.0 https://github.com/ultralytics/yolov5.git
进入yolo目录
cd yolov5
下载权重:
wget https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt
这将下载一个yolov5.pt
文件
转换模型格式:YOLOv5默认使用PyTorch框架。为了使用TensorRT优化,你需要将PyTorch模型转换成ONNX格式。使用以下命令可以完成这一转换:
python export.py --weights yolov5s.pt --img 640 --batch 1 --device 0 --opset 12 --simplify --include onnx
这将生成一个.onnx
文件,例如yolov5s.onnx
。
在yolo5文件目录下,在终端下载依赖项:
pip3 install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
步骤1所有操作可以在其它设备上进行
步骤 2: 使用TensorRT优化模型
安装TensorRT:确保你的Jetson Nano上安装了TensorRT。可以通过JetPack SDK manager安装。
git clone -b yolov5-v7.0 https://github.com/warng-xinyu/tensorrtx.git
回到yolov5目录
cd ~/yolov5
复制gen_wts.py到yolov5
cp ~/tensorrtx/yolov5/gen_wts.py .
生成wts文件
python3 gen_wts.py -w yolov5s.pt -o yolov5s.wts
步骤2以上所有操作可以在其它设备上进行,但下面操作要在最终运行的设备jetson nano上进行
进入tensorrtx目录
cd ~/tensorrtx
创建build文件夹,并进入
mkdir build
cd build
复制之前生成的wts文件复制到build目录下
cp ~/yolov5/yolov5s.wts .
执行cmake命令,然后执行make进行c++的编译
cmake ..
make
过程中可能报一些警告,可以忽略
构建TensorRT引擎:使用trtexec
工具或编写自定义代码来从ONNX模型构建TensorRT优化引擎。例如:
方式一(推荐):使用自定义的yolov5_det工具,从yolov5s.wts的权重文件和yolov5s.engine的引擎文件生成了另一个引擎文件
./yolov5_det -s yolov5s.wts yolov5s.engine s
方式二:令使用了trtexec工具,从yolov5s.onnx的ONNX模型文件生成了yolov5s.engine的TensorRT引擎文件
/usr/src/tensorrt/bin/trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s.engine --explicitBatch --fp16
explicitBatch指定了显式批处理,–fp16表示使用16位浮点精度进行计算
这将创建一个TensorRT引擎文件yolov5s.engine
检验识别效果
./yolov5_det -d yolov5s.engine ../images
步骤 3: 配置DeepStream
安装DeepStream:确保你的Jetson Nano上安装了DeepStream SDK。
准备配置文件:修改DeepStream的配置文件以使用YOLOv5模型。你需要更新配置文件中的模型文件路径、输入尺寸、模型输出解析等。确保模型的输出层名称、类别数等与YOLOv5模型相匹配。
YOLO Layer Plugin:因为YOLOv5的输出可能需要自定义的解析器,你可能需要实现一个自定义插件来解析网络输出。这通常涉及到编写一个继承自nvdsinfer
的C++类,并在DeepStream的配置中指定这个插件。
步骤 4: 运行DeepStream应用
启动应用:使用配置好的DeepStream管道启动应用。这将自动管理视频流的输入、YOLOv5模型的推理及结果的输出和显示。
deepstream-app -c deepstream_app_config.txt
注意事项
在部署过程中,确保所有的路径和配置文件设置正确无误。
Jetson Nano的资源有限,使用较小的模型和开启精度减少(如FP16)可能会帮助提高性能。
通过以上步骤,你可以在Jetson Nano上使用TensorRT和DeepStream成功部署和加速YOLOv5算法。这样做可以显著提高处理速度,特别是在处理高分辨率或多视频流输入时。