当前位置: 首页 > 工具软件 > nano > 使用案例 >

Yolov5—nano部署

楚良平
2023-12-01

一、前言

1.使用YOLOv5-4x版本在服务器端进行模型训练
2.使用YOLOv5-4x版本在nano端进行模型部署

二、服务器端

1.训练模型保存权重文件(.pt)

2.https://github.com/ultralytics/yolov5

#clone yolov5项目选择下载V4.0版本的源码
#若提示Command 'git' not found,则执行:
sudo apt install git
git clone https://github.com/ultralytics/yolov5

3.https://github.com/wang-xinyu/tensorrtx

#clone tensorrt项目并下载
git clone https://github.com/wang-xinyu/tensorrtx

4.获得.wts文件(.pt—>.wts)

将tensorrtx/yolov5/gen_wts.py复制到ultralytics/yolov5,并修改gen_wts.py文件中权重文件的名称和所要生成wts文件的名称。

三、nano端

1.预备工作
#使用tensorrt进行量化加速,生成engine文件
https://github.com/wang-xinyu/tensorrtx

#clone tensorrtx
git clone https://github.com/wang-xinyu/tensorrtx

1.1修改配置参数

1.1.1 yolov5.cpp

#define USE_FP16
#define DEVICE 0
#define NMS_THRESH 0.4
#define CONF_THRESH 0.5
#define BATCH_SIZE 6
#define NET s

其中BATCH_SIZE参数影响的是模型可以实现多路输出,NET参数表示所用yolov5权重的类型(s、l、m、x)

1.1.2yololayer.h

CLASS_NUM = 10

该参数表示可检测物体的类别数,按个人需求进行更改

2.编译tensorrtx/yolov5

cd /tensorrtx/yolov5
mkdir build
cd build/
cmake ..
make

3.生成engine文件
copy yolov5.wts 文件到tensorrtx/yolov5/build目录下

./yolov5的参数设置:

./yolov5 -s [.wts] [.engine] [s/m/l/x/s6/m6/l6/x6 or c/c6 gd gw]  // serialize model to plan file
./yolov5 -d [.engine] ../samples  // deserialize plan file and run inference

所以执行下述指令生成对应的.engine文件

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

注意:需要在对应的使用设备上生成engine文件。因为engine文件是和平台相关的,当换为另一台nano时,需要在重新进行生成

4.部署(使用DeepStream部署yolov5s)

#准备阶段,下载DeepStream-Yolo项目
https://github.com/marcoslucianops/DeepStream-Yolo
选择external文件夹下的yolov5-4.0,并将其copy到/opt/nvidia/deepstream/seepstream/sources/目录下

4.1 修改配置参数

#修改可检测目标的类别数量
config-infer-primary.txt:num_detected_classes = 1

#修改类别标签
labels.txt: 写入所要检测目标的标签

#进入nvdsinfer_custom_impl_Yolo文件夹
#修改极大值抑制参数
nvdsparsebbox_Yolo.cpp:
#define KNMS_THRESH 0.45

#修改yololayer.h,定义为自己所需要检测的类别数量
static constexpr int CLASS_NUM = 80

其中 KNMS_THRESH参数表示,当多个检测框重合率超过45%时,就判定为一个物体.
当检测目标比较密集,且希望每个目标都能用区域框标出来时,可以将数值增大。

4.2编译

cd /opt/nvidia/deepstream/deepstream/sources/yolov5-4.0
CUDA_VER=10.2 make -C nvdsinfer_custom_impl_Yolo
#编译后生成.so文件,即动态链接库文件
libnvdsinfer_custom_impl_Yolo.so
 类似资料: