Skip to content

Latest commit

 

History

History

tensorrt_yolov4

Folders and files

NameName
Last commit message
Last commit date

parent directory

..
data
data
 
 
source
source
 
 
Makefile
Makefile
 
 
Makefile.config
Makefile.config
 
 
README.md
README.md
 
 

README.md

YOLOv4 Standalone Program of Multi-Tasks

1. Contents

  • common Some common code dependencies and utilities

  • source Source code of standalone Program

    • main.cpp: Program main entrance where parameters are configured here
    • SampleYolo.hpp: YOLOv4 inference class definition file
    • SampleYolo.cpp: YOLOv4 inference class functions definition file
    • onnx_add_nms_plugin.py: Python script to add BatchedNMSPlugin node into ONNX model
    • generate_coco_image_list.py: Python script to get list of image names from MS COCO annotation or information file

  • data This directory saves:

    • yolov4.onnx: the ONNX model (User generated)
    • yolov4.engine: the TensorRT engine model (would be generated by this program)
    • demo.jpg: The demo image (Already exists)
    • demo_out.jpg: Image detection output of the demo image (Already exists, but would be renewed by the program)
    • names.txt: MS COCO dataset label names (have to be downloaded or generated via COCO API)
    • categories.txt: MS COCO dataset categories where IDs and names are separated by "\t" (have to be generated via COCO API)
    • val2017.txt: MS COCO validation set image list (have to be generated from corresponding COCO annotation file)
    • valdev2017.txt: MS COCO test set image list (have to be generated from corresponding COCO annotation file)
    • coco_result.json: MS COCO dataset output (would be generated by this program)

2 Prerequisites before building & running YOLOv4 standalone

2.1 Download TensorRT (higher than 7.1, you can ignore this step if TensorRT 7.1 is already installed)

2.2 Download and build TensorRT OSS

2.3 Generate YOLOv4 ONNX model with BatchedNMSPlugin node included

Step 1 Generate YOLOv4 ONNX model (CSPDarknet-53 CNN + YOLO header CNN + YOLO layers)

Step 2 Add into YOLOv4 ONNX model the BatchedNMSPlugin (CSPDarknet-53 CNN + YOLO header CNN + YOLO layers + BatchedNMSPlugin)

How can I add BatchedNMSPlugin node into ONNX model?

  • Open source_gpu_nms/onnx_add_nms_plugin.py

  • Update attribute values to suit your model

Example:

    attrs["shareLocation"] = 1
    attrs["backgroundLabelId"] = -1
    attrs["numClasses"] = 80
    attrs["topK"] = topK # from program arguments
    attrs["keepTopK"] = keepTopK # from program arguments
    attrs["scoreThreshold"] = 0.3
    attrs["iouThreshold"] = 0.6
    attrs["isNormalized"] = 1
    attrs["clipBoxes"] = 1

  • Copy onnx_add_nms_plugin.py into <TensorRT_OSS_dir>/tools/onnx-graphsurgeon

  • Go to <TensorRT_OSS_dir>/tools/onnx-graphsurgeon and execute onnx_add_nms_plugin.py

cd <TensorRT_OSS_dir>/tools/onnx-graphsurgeon
python onnx_add_nms_plugin.py -f <yolov4_onnx_file> -t <topk_value> -k <keep_topk_value>

3. How can I build and run YOLOv4 standalone program?

3.1 Add common source code includes

  • This YOLOv4 standalone sample depends on the same common includes as other C++ samples of TensorRT.
    • Option 1: Add a link to <where_tensorRT_is_installed>/TensorRT-7.1.x.x/samples/common in tensorrt_yolov4
    cd <dir_on_your_machine>/yolov4_sample/tensorrt_yolov4
ln -s <where_tensorRT_is_installed>/TensorRT-7.1.x.x/samples/common common
    • Option 2: Simply copy common includes into tensorrt_yolov4
    cd <dir_on_your_machine>/yolov4_sample/tensorrt_yolov4
cp -r <where_tensorRT_is_installed>/TensorRT-7.1.x.x/samples/common common ./

3.2 OpenCV dependencies

  • Note: There are OpenCV dependencies in this program. Please check if there are OpenCV includes in /usr/include/opencv and if OpenCV libraries like -lopencv_core and -lopencv_imgproc are installed.

  • Follow README and documents of this repository https://github.com/opencv/opencv to install OpenCV if corresponding includes and libraries do not exist.

3.3 Compile and build

cd <dir_on_your_machine>/yolov4_sample/yolo_cpp_standalone/source_gpu_nms
make clean
make -j<num_processors>

3.4 Basic program parameters

  • Step1: Use text editor to open main.cpp in <dir_on_your_machine>/YOLOv4_Sample/tensorrt_yolov4/source

  • Step2: Go to where function initializeSampleParams() is defined

  • Step3: You will find some basic configurations in initializeSampleParams() like follows:

    // This argument is for calibration of int8
    // Int8 calibration is not available until now
    // You have to prepare samples for int8 calibration by yourself 
    params.nbCalBatches = 80;

    // The engine file to generate or to load
    // The engine file does not exist:
    //     This program will try to load onnx file and convert onnx into engine
    // The engine file exists:
    //     This program will load the engine file directly
    params.engingFileName = "../data/yolov4.engine";

    // The onnx file to load
    params.onnxFileName = "../data/yolov4.onnx";
    
    // Input tensor name of ONNX file & engine file
    params.inputTensorNames.push_back("input");
    
    // Old batch configuration, it is zero if explicitBatch flag is true for the tensorrt engine
    // May be deprecated in the future
    params.batchSize = 0;
    
    // Number of classes (usually 80, but can be other values)
    params.outputClsSize = 80;
    
    // topK parameter of BatchedNMSPlugin
    params.topK = 2000;
    
    // keepTopK parameter of BatchedNMSPlugin
    params.keepTopK = 1000;

    // Batch size, you can modify to other batch size values if needed
    params.explicitBatchSize = 1;

    params.inputImageName = "../data/demo.jpg";
    params.cocoClassNamesFileName = "../data/coco.names";
    params.cocoClassIDFileName = "../data/categories.txt";

    // Config number of DLA cores, -1 if there is no DLA core
    params.dlaCore = -1;
  • Step4: Copy and rename the ONNX file (BatchedNMSPlugin node included) to the location defined by initializeSampleParams()

3.5 Run this program to convert ONNX file into Engine file

  • This program will automatically convert ONNX into engine if engine does not exist.
  • Command:
    • To generate Engine of fp32 mode:
    ../bin/yolov4
    • To generate Engine of fp16 mode:
    ../bin/yolov4 --fp16

3.6 Specific program parameters for demo mode, speed mode and coco mode

3.6.1 To run this program in demo mode

  • Command:
../bin/yolov4 --demo
  • This program will feed the demo image into YOLOv4 engine and write detection output as an image.
  • Please make sure params.demo = 1 if you want to run this program in demo mode.
    // Configurations to run a demo image
    params.demo = 1;
    params.outputImageName = "../data/demo_out.jpg";

3.6.2 To run this program in speed mode

  • Command:
../bin/yolov4 --speed
  • This program will repeatedly feed the demo image into engine to accumulate time consumed in each iteration
  • Please make sure params.speedTest = 1 if you want to run this program in speed mode
    // Configurations to run speed test
    params.speedTest = 1;
    params.speedTestItrs = 1000;

3.6.3 To run this program in coco mode

  • Command:
../bin/yolov4 --coco
  • Corresponding configuration in initializeSampleParams() would be like this:
    // Configurations of Test on COCO dataset
    params.cocoTest = 1;
    params.cocoClassNamesFileName = "../data/coco.names";
    params.cocoClassIDFileName = "../data/categories.txt";
    params.cocoImageListFileName = "../data/val2017.txt";
    params.cocoTestResultFileName = "../data/coco_result.json";
    params.cocoImageDir = "../data/val2017";

Note: COCO dataset is just an example, you can use your own validation set or test set to validate YOLOv4 model trained by your own training set

python generate_coco_image_list.py <json file of image annotations> <image list text>

  • For example, to generate validation image list, the command would be:

    python generate_coco_image_list.py instances_val2017.json val2017.txt

  • For example, to generate test-dev image list, the command would be:

    python generate_coco_image_list.py image_info_test-dev2017.json testdev2017.txt

  • This program will read image names from the list file whose path should be the same as params.cocoImageListFileName, and then feed these images located in params.cocoImageDir to YOLOv4 engine

  • Please make sure params.cocoTest = 1 and images exist in params.cocoImageDir