spacewalk01 / tensorrt-yolov6

C++ implementation of YOLOv6 using TensorRT

TensorRT-YOLOv6

This repo provides the C++ implementation of YOLOv6 based on TensorRT for performing object detection in real-time.

Environment

• Windows 10
• Visual Studio 2017
• RTX 3070 GPU
• TensorRT 7.2.1.6
• CUDA 11.1, Cudnn 8
• OpenCV 4.5.1 with CUDA

Benchmark

The following result was obtained after performing inference 1000 iterations. It includes image pre-processing, model inference and post-processing (non-max suppression) times. The resolution of the test video was 1280x720.

Model	Device	Mode	Input Shape(HxW)	FPS
YOLOv6-n	RTX3070	FP16	640x640	214
YOLOv6-n	RTX3070	FP32	640x640	165
YOLOv6-tiny	RTX3070	FP16	640x640	193
YOLOv6-tiny	RTX3070	FP32	640x640	128
YOLOv6-s	RTX3070	FP16	640x640	130
YOLOv6-s	RTX3070	FP32	640x640	115

Tutorial

• Step1: Download a pretrained onnx model from the YOLOv6 release

• Step2: Open the solution with Visual Studio and select x64 and Release for the configuration

• Step3: Go to where function initializeSampleParams() is defined and modify the parameters

    // 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 = "yolov6s.engine";

    // The onnx file to load
    params.onnxFileName = "yolov6s.onnx";
    
    // Number of classes (usually 80, but can be other values)
    params.outputClsSize = 80;
    
    // Input video filename
    params.inputVideoName = "test.mp4";
    // Model class filename
    params.cocoClassNamesFileName = "coco.names";
    
    // Threshold values
    params.confThreshold = 0.3;
    params.nmsThreshold = 0.5;
    
    // Floating point precision
    //     By default, fp32
    params.fp16 = true; 

Todo

• Benchmark Test with different yolov6 versions
• Documentation
• Fix arguments

Reference

• yolov4_deepstream
• yolov6