git clone https://github.com/meituan/YOLOv6.git
python deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --batch 1
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 |
Here is a Python Demo mybe help quickly understand this repo Link
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>