deepstream_yolo/YOLOv5-3.X.md

YOLOv5

NVIDIA DeepStream SDK 5.1 configuration for YOLOv5 3.0/3.1 models

Thanks DanaHan, wang-xinyu and Ultralytics

• Requirements
• Convert PyTorch model to wts file
• Convert wts file to TensorRT model
• Compile nvdsinfer_custom_impl_Yolo
• Testing model

Requirements

• TensorRTX

• Ultralytics

• Matplotlib (for Jetson plataform)

sudo apt-get install python3-matplotlib

• PyTorch (for Jetson plataform)

wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev
pip3 install Cython
pip3 install numpy torch-1.8.0-cp36-cp36m-linux_aarch64.whl

• TorchVision (for Jetson platform)

sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
git clone --branch 0.9.0 https://github.com/pytorch/vision torchvision
cd torchvision
export BUILD_VERSION=0.9.0  
python3 setup.py install --user

Convert PyTorch model to wts file

1. Download repositories

git clone https://github.com/DanaHan/Yolov5-in-Deepstream-5.0.git yolov5converter
git clone -b yolov5-v3.1 https://github.com/wang-xinyu/tensorrtx.git
git clone -b v3.1 https://github.com/ultralytics/yolov5.git

2. Download latest YoloV5 (YOLOv5s, YOLOv5m, YOLOv5l or YOLOv5x) weights to yolov5/weights directory (example for YOLOv5s)

wget https://github.com/ultralytics/yolov5/releases/download/v3.1/yolov5s.pt -P yolov5/weights/

3. Copy gen_wts.py file (from tensorrtx/yolov5 folder) to yolov5 (ultralytics) folder

cp tensorrtx/yolov5/gen_wts.py yolov5/gen_wts.py

4. Generate wts file

cd yolov5
python3 gen_wts.py

yolov5s.wts file will be generated in yolov5 folder

Note: if you want to generate wts file to another YOLOv5 model (YOLOv5m, YOLOv5l or YOLOv5x), edit get_wts.py file changing yolov5s to your model name

model = torch.load('weights/yolov5s.pt', map_location=device)['model'].float()  # load to FP32
model.to(device).eval()

f = open('yolov5s.wts', 'w')

Convert wts file to TensorRT model

1. Replace yololayer files from tensorrtx/yolov5 folder to yololayer and hardswish files from yolov5converter

mv yolov5converter/yololayer.cu tensorrtx/yolov5/yololayer.cu
mv yolov5converter/yololayer.h tensorrtx/yolov5/yololayer.h

2. Move generated yolov5s.wts file to tensorrtx/yolov5 folder (example for YOLOv5s)

cp yolov5/yolov5s.wts tensorrtx/yolov5/yolov5s.wts

3. Build tensorrtx/yolov5

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

4. Convert to TensorRT model (yolov5s.engine file will be generated in tensorrtx/yolov5/build folder)

sudo ./yolov5 -s

5. Create a custom yolo folder and copy generated files (example for YOLOv5s)

mkdir /opt/nvidia/deepstream/deepstream-5.1/sources/yolo
cp yolov5s.engine /opt/nvidia/deepstream/deepstream-5.1/sources/yolo/yolov5s.engine

Note: by default, yolov5 script generate model with batch size = 1, FP16 mode and s model.

#define USE_FP16  // comment out this if want to use FP32
#define DEVICE 0  // GPU id
#define NMS_THRESH 0.4
#define CONF_THRESH 0.5
#define BATCH_SIZE 1

#define NET s  // s m l x

Edit yolov5.cpp file before compile if you want to change this parameters.

Compile nvdsinfer_custom_impl_Yolo

1. Run command

sudo chmod -R 777 /opt/nvidia/deepstream/deepstream-5.1/sources/

2. Donwload my external/yolov5 folder and move files to created yolo folder

3. Compile lib

• x86 platform

cd /opt/nvidia/deepstream/deepstream-5.1/sources/yolo
CUDA_VER=11.1 make -C nvdsinfer_custom_impl_Yolo

• Jetson platform

cd /opt/nvidia/deepstream/deepstream-5.1/sources/yolo
CUDA_VER=10.2 make -C nvdsinfer_custom_impl_Yolo

Testing model

Use my edited deepstream_app_config.txt and config_infer_primary.txt files available in my external/yolov5-3.X folder

Run command

deepstream-app -c deepstream_app_config.txt

Note: based on selected model, edit config_infer_primary.txt file

For example, if you using YOLOv5x

model-engine-file=yolov5s.engine

to

model-engine-file=yolov5x.engine

To change NMS_THRESH, edit nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp file and recompile

#define kNMS_THRESH 0.45

To change CONF_THRESH, edit config_infer_primary.txt file

[class-attrs-all]
pre-cluster-threshold=0.25