Skip to content

Latest commit

 

History

History
72 lines (50 loc) · 1.66 KB

README.md

File metadata and controls

72 lines (50 loc) · 1.66 KB

YOLOV6-TensorRT in C++

Step 1: 准备TRT序列化引擎

git clone https://github.com/meituan/YOLOv6.git

导出onnx

python deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --batch 1

转化为TensorRT Engine

python export.py -o onnx-name -e trt-name -p fp32/16/int8

测试

cd yolov6
python trt.py
model input FPS Device Language
yolov6s 640*640 FP16 360FPS A100 Python
yolov6s 640*640 FP32 350FPS A100 Python
yolov6s 640*640 FP32 330FPS 1080Ti C++
yolov6s 640*640 FP32 300FPS 1080Ti Python

bilibili

Here is a Python Demo mybe help quickly understand this repo Link

Step 2: C++

Please follow the TensorRT Installation Guide to install TensorRT.

And you should set the TensorRT path and CUDA path in CMakeLists.txt.

If you train your custom dataset, you may need to modify the value of num_class.

const int num_class = 80;

Install opencv with sudo apt-get install libopencv-dev (we don't need a higher version of opencv like v3.3+).

build the demo:

mkdir build
cd build
cmake ..
make

Then run the demo:

./yolov6 ../model_trt.engine -i ../../../../assets/dog.jpg

or

./yolov6 <path/to/your/engine_file> -i <path/to/image>