DeepStream5.0部署YOLOv5(AGX)

00 2

已于 2022-06-15 17:00:10 修改

阅读量2.4k

点赞数

文章标签：深度学习 pytorch 目标检测

于 2022-02-27 09:09:39 首次发布

本文链接：https://blog.csdn.net/weixin_41170972/article/details/120515631

版权

转载:

好耶

背景

本文旨在对 deepstream 中使用 yolov5 的方法做一介绍
测试环境： Ubuntu 18.04, CUDA 10.2, AGX, jetpack4.5.1

deepstream-app --version-all

请添加图片描述

jtop

请添加图片描述

yolo版本：yolov5的3.1version (我用version5.0的时候有问题)
项目地址：https://github.com/DanaHan/Yolov5-in-Deepstream-5.0

PYTHON环境

之前写过:AGX部署yolov5(pytorch)

源码

mkdir deepstream_yolov5
cd deepstream_yolov5
git clone https://github.com/DanaHan/Yolov5-in-Deepstream-5.0.git
git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/ultralytics/yolov5.git

下载预训练模型

在 yolov5 界面下载预训练模型（3.1版本模型），这里我们以 YOLOv5s 为例介绍
地址：https://github.com/ultralytics/yolov5
下载好模型后放到 yolov5/weights 目录下

生成engine

pip install scikit-build
conda activate py36
cd yolov5
cp ../tensorrtx/yolov5/gen_wts.py ./

修改gen_wts.py的pt_file, wts_file路径

import sys
import argparse
import os
import struct
import torch
from utils.torch_utils import select_device

pt_file, wts_file = 'weights/yolov5s.pt','yolov5s.wts'
# Initialize
device = select_device('cpu')
# Load model
model = torch.load(pt_file, map_location=device)['model'].float()  # load to FP32
model.to(device).eval()

with open(wts_file, 'w') as f:
    f.write('{}\n'.format(len(model.state_dict().keys())))
    for k, v in model.state_dict().items():
        vr = v.reshape(-1).cpu().numpy()
        f.write('{} {} '.format(k, len(vr)))
        for vv in vr:
            f.write(' ')
            f.write(struct.pack('>f' ,float(vv)).hex())
        f.write('\n')

运行后会在本地生成 yolov5x.wts 文件，然后拷贝到 Yolov5-in-Deepstream-5.0 文件夹中

cp yolov5x.wts ../Yolov5-in-Deepstream-5.0
cd ../Yolov5-in-Deepstream-5.0

修改 yolov5.cpp 文件，将 NET 宏改成自己对应的模型

#define NET x // s m l x
然后编译

mkdir build
cd build
cmake ..
make
sudo ./yolov5 -s

如果要用自己训练的权重，要改yololayer.h

检查

运行后会生成 yolov5x.engine 和 libmyplugin.so 文件，首先检查下engine运行结果是否正确
创建测试图片

mkdir ../samples
#网上下载两张 coco 数据集图片
sudo ./yolov5 -d  ../samples

请添加图片描述
可以在build目录下生成检测结果图片，格式是_name.jpg，打开图片可以看到检测框正常
验证完毕后，将生成的模型和动态库拷贝到 deepstream 目录中

deepstream部署

1.在 Deepstream 5.0/nvdsinfer_custom_impl_Yolo 目录中运行 make 编译，生成 libnvdsinfer_custom_impl_Yolo.so 文件

2.返回上一级，修改 config_infer_primary_yoloV5.txt 文件，总共有以下几处需要修改
model-engine-file=yolov5s.engine --> model-engine-file=yolov5s.engine
custom-lib-path = objectDetector_Yolo_V5/nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so --> custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so
原工程中缺少 labels.txt 文件，把 coco 数据集的 80 类 labels 拷贝过来即可

然后根据自己的需求，修改 deepstream_app_config_yoloV5.txt 即可

然后载入定制化的动态库，并运行即可

LD_PRELOAD=./libmyplugins.so deepstream-app -c deepstream_app_config_yoloV5.txt

会重新生成 engine 文件，并可以看到运行结果
注意，engine 文件要在 Yolov5-in-Deepstream-5.0 工程中生成，在 tensorrtx 工程中生成 engine 再导过来运行会出现一大堆框
请添加图片描述
labels.txt

1,1,person
2,2,bicycle
3,3,car
4,4,motorcycle
5,5,airplane
6,6,bus
7,7,train
8,8,truck
9,9,boat
10,10,traffic_light
11,11,fire_hydrant
13,12,stop_sign
14,13,parking_meter
15,14,bench
16,15,bird
17,16,cat
18,17,dog
19,18,horse
20,19,sheep
21,20,cow
22,21,elephant
23,22,bear
24,23,zebra
25,24,giraffe
27,25,backpack
28,26,umbrella
31,27,handbag
32,28,tie
33,29,suitcase
34,30,frisbee
35,31,skis
36,32,snowboard
37,33,sports_ball
38,34,kite
39,35,baseball_bat
40,36,baseball_glove
41,37,skateboard
42,38,surfboard
43,39,tennis_racket
44,40,bottle
46,41,wine_glass
47,42,cup
48,43,fork
49,44,knife
50,45,spoon
51,46,bowl
52,47,banana
53,48,apple
54,49,sandwich
55,50,orange
56,51,broccoli
57,52,carrot
58,53,hot_dog
59,54,pizza
60,55,donut
61,56,cake
62,57,chair
63,58,couch
64,59,potted_plant
65,60,bed
67,61,dining_table
70,62,toilet
72,63,tv
73,64,laptop
74,65,mouse
75,66,remote
76,67,keyboard
77,68,cell_phone
78,69,microwave
79,70,oven
80,71,toaster
81,72,sink
82,73,refrigerator
84,74,book
85,75,clock
86,76,vase
87,77,scissors
88,78,teddy_bear
89,79,hair_drier
90,80,toothbrush