点击查看系列文章目录
0 背景
deepstream 5.0 GA版本的 SDK 中包含有 yolov3 的检测,还没有正式加上 yolov4 的功能,不过 nvidia 官方也在 github 上发布了相关的代码,本文对使用方法做一个介绍
1 模型准备
官方没有提供转换后的模型,需要我们自己转换
git clone https://github.com/Tianxiaomo/pytorch-YOLOv4.git
cd pytorch-YOLOv4
建议创建 conda 环境来安装部署相关软件
conda create -n yolov4 python=3.6
conda activate yolov4
pip install -r requirement.txt
# conda install -c menpo opencv
pip install protobuf
sudo apt-get install protobuf-compiler libprotoc-dev
pip install onnx
pip install onnxruntime
注意:这里要安装 requirement 里边的软件版本,尤其是 pytorch 要安装1.4.0版本,否则其它版本会有各种问题,其中一个问题就是生成 onnx 模型时报错如下
[12/22/2020-17:35:18] [E] [TRT] Layer: Where_377's output can not be used as shape tensor.
[12/22/2020-17:35:18] [E] [TRT] Network validation failed.
[12/22/2020-17:35:18] [E] Engine creation failed
[12/22/2020-17:35:18] [E] Engine set up failed
如果是jetson上,torch的安装方法参考《Jetson之刷机后配置大全》2.4章节上边
下载 darknet 版本的 预训练模型和配置文件,放到 pytorch-YOLOv4 路径下
yolov4.cfg文件支持不同尺度的输入,对应的精度和计算量不一样,根据自己的需求选择一个长宽输入
然后使用脚本转化为 ONNX 模型
python demo_darknet2onnx.py yolov4.cfg yolov4.weights ./data/dog.jpg 1
经过上述指令,会生成 *.onnx 模型,然后再利用 trtexec 工具生成 engine 模型
trtexec --onnx=yolov4_1_3_512_320.onnx --explicitBatch --saveEngine=yolov4_1_3_320_512_fp16.engine --workspace=4096 --fp16
这个过程比较耗时,根据自己 GPU 性能决定,我在 T4 显卡上用了大概 10 分钟左右生成,在Xavier AGX用了20分钟
转化过程中有 warning 如下,可以不用管
[W] [TRT] onnx2trt_utils.cpp:198: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
有的人可能找不到 trtexec 文件,这个是在tensorrt的安装包里,可以使用find指令查找一下
sudo find / -name trtexec
注意:生成的engine是平台相关的,在x86上生成的engine不能在jetson上使用,需要使用上边的方法重新在jetson上生成一遍
2 代码准备
找到 deepstream sdk 的安装路径,默认为 /opt/nvidia/deepstream/deepstream-5.0,进入 source 路径,克隆代码
git clone https://github.com/NVIDIA-AI-IOT/yolov4_deepstream.git
cd yolov4_deepstream/deepstream_yolov4/nvdsinfer_custom_impl_Yolo
export CUDA_VER=10.2
make
编译成功后,生成 libnvdsinfer_custom_impl_Yolo.so 动态库,如果编译报错如下
nvdsparsebbox_Yolo.cpp:24:10: fatal error: nvdsinfer_custom_impl.h: No such file or directory
#include "nvdsinfer_custom_impl.h"
^~~~~~~~~~~~~~~~~~~~~~~~~
compilation terminated.
Makefile:41: recipe for target 'nvdsparsebbox_Yolo.o' failed
make: *** [nvdsparsebbox_Yolo.o] Error 1
是因为 makefile 中 includes 路径写错了,把第25行改成如下
-- CFLAGS+= -I../../includes -I/usr/local/cuda-$(CUDA_VER)/include
++ CFLAGS+= -I../../../includes -I/usr/local/cuda-$(CUDA_VER)/include
代码提供两种方式运行 yolov4,分别是直接使用 TensorRT 推理,另一种方式是使用 deepstream 做推理,我们主要关心在 deepstream 中的使用方法
3 修改配置文件
修改 deepstream_app_config_yoloV4.txt 和 config_infer_primary_yoloV4.txt 中的 model-engine-file 为第一步生成的 engine 文件,修改自己的 source 和 sink,然后就可以使用 deepstream-app 检测,我这里使用samples/streams/sample_1080p_h264.mp4 做测试,
deepstream-app -c deepstream_app_config_yoloV4.txt
检测效果如下
4 集成自定义模型
如果我们使用 darknet 训练了自己的 yolov4 模型,集成到 deepstream 中时需要重新生成 libnvdsinfer_custom_impl_Yolo.so 动态库,文件中有几个地方要修改:
- nvdsparsebbox_Yolo.cpp 中 26 行 static const int NUM_CLASSES_YOLO = 80 改为自己的类别数
- deepstream_yolov4/labels.txt 修改为自己的类别
修改完之后重新make即可,其它地方一样