在deepstream上使用yolov4的模型
目前nvidia官方并没有开发出适配Yolov4的配置文件和对应的代码,但是现在已经有第三方的解决方案:
首先假设你已经通过Alex AB训练得到了一个权重文件(weightFile)yolov4.weights,并且还有一个对应的配置文件(cfgFile)yolov4.cfg
生成TRT模型
执行
- git clone https://github.com/Tianxiaomo/pytorch-YOLOv4.git
- cd pytorch-YOLOv4
Requirement
pytorch >= 1.4
onnx
onnxruntime
opencv-python
Pytorch version Recommended:
- Pytorch 1.4.0 for TensorRT 7.0 and higher
- Pytorch 1.5.0 and 1.6.0 for TensorRT 7.1.2 and higher
当所有要求的库安装之后,在项目目录下执行。
python demo_darknet2onnx.py <cfgFile> <weightFile> <imageFile> <batchSize>
这个命令会生成两个模型文件(.onnx),一个默认batchsize为1用来进行demo的运行,一个是你设置的batchsize。
在你的deepstream平台上,把你需要的.onnx文件复制过去,执行
trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
说明:这里的trtexec是在你安装完TensorRT时候有的,如果直接运行提示没有这个指令,那么使用绝对路径,位置为:/usr/src/tensorrt/bin/trtexec
这时候onnx选项填你自己的.onnx模型文件,saveEngine填你要输出的模型文件名,实测–workspace可以忽略,貌似默认值是16。下面是我的指令
- /usr/src/tensorrt/bin/trtexec --onnx=yolov4-fire.onnx --explicitBatch --saveEngine=yolov4-fire.engine --fp16
经过漫长的等待之后,你可以获得一个.engine文件。
在deepstream上使用TRT模型
前面提到过,nvidia团队并没有对yolov4开发相应的app,所以有些需要我们自己修改。在这里我在deepstream环境中创建了一个文件夹:Yolov4-fire
在文件夹下面我把.engine文件复制了过去。
接下来!!!!!把github项目中DeepStream/nvdsinfer_custom_impl_Yolo 文件夹复制到Yolov4-fire文件夹下。修改nvdsparsebbox_Yolo.cpp中的NUM_CLASSES_YOLO为你自己的类别数,执行
- $ export CUDA_VER=10.2
- $ make -C nvdsinfer_custom_impl_Yolo
然后根据需要添加并且修改两个配置文件config_infer_primary_yoloV4.txt、deepstream_app_config_yoloV4.txt和label文件,它们都在github项目下DeepStream文件夹中。
你需要修改的是:
- model-engine-file
- labelfile-path
- 配置文件中的输入输出配置
最后在目录下执行
- deepstream-app -c deepstream_app_config_yoloV4.txt
大功告成,App run successfully