TensorRT-For-YOLO-Series项目：实现yolov10模型的python-tensorrt推理（对比int8与fp16推理差异）

项目地址：https://github.com/Linaom1214/TensorRT-For-YOLO-Series/tree/cuda-python
算法支持状态：
2024.6.16 Support YOLOv9, YOLOv10, changing the TensorRT version to 10.0
2023.8.15 Support cuda-python
2023.5.12 Update
2023.1.7 support YOLOv8
2022.11.29 fix some bug thanks @JiaPai12138
2022.8.13 rename reop、 public new version、 C++ for end2end
2022.8.11 nms plugin support ==> Now you can set --end2end flag while use export.py get a engine file
2022.7.8 support YOLOv7
2022.7.3 support TRT int8 post-training quantization

1、tensorrt环境安装

下载tensorrt https://developer.nvidia.com/tensorrt/download/10x
这里根据个人的系统、cuda版本进行选择。

解压文件，并将tensorrt的路径添加系统环境变量中，具体如下所示：
然后安装python-tensorrt，具体如下所示

最后安装
pip install cuda-python

同时需要设计cuda环境变量
1、在系统环境变量path中添加 %CUDA_PATH%

2、在系统环境变量中添加 CUDA_PATH

2、onnx2trt

执行以下代码可以将pt模型导出为onnx模型

from ultralytics import YOLO

model = YOLO("yolov10n.pt")
model.fuse()  
model.info(verbose=False)  # Print model information
model.export(format='onnx')  # TODO: 

执行以下命令可以将onnx模型转换为trt模型
python export.py -o yolov10n.onnx -e yolov10.trt --end2end --v10 -p fp32

3、模型推理

测试代码如下所示

python trt.py -e yolov10.trt  -i src/1.jpg -o yolov10-1.jpg --end2end

推理输出如下所示，针对640x640的输入，fps为146，基本上是6.8ms一张图。推理软硬件环境为：win10+tensorrt10.4+python3.8+cuda12.1+3060显卡

Namespace(end2end=True, engine='yolov10.trt', image='src/1.jpg', output='yolov10-1.jpg', video=None)
146.50372522347308 FPS

对应的trt.py代码如下所示

from utils.utils import preproc, vis
from utils.utils import BaseEngine
import numpy as np
import cv2
import time
import os
import argparse

class Predictor(BaseEngine):
    def __init__(self, engine_path):
        super(Predictor, self).__init__(engine_path)
        self.n_classes = 80  # your model classes

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-e", "--engine", help="TRT engine Path")
    parser.add_argument("-i", "--image", help="image path")
    parser.add_argument("-o", "--output", help="image output path")
    parser.add_argument("-v", "--video",  help="video path or camera index ")
    parser.add_argument("--end2end", default=False, action="store_true",
                        help="use end2end engine")

    args = parser.parse_args()
    print(args)

    pred = Predictor(engine_path=args.engine)
    pred.get_fps()
    img_path = args.image
    video = args.video
    if img_path:
      origin_img = pred.inference(img_path, conf=0.1, end2end=args.end2end)

      cv2.imwrite("%s" %args.output , origin_img)
    if video:
      pred.detect_video(video, conf=0.1, end2end=args.end2end) # set 0 use a webcam

4、fp16推理

导出fp16模型

python export.py -o yolov10n.onnx -e yolov10.trt --end2end --v10 -p fp16
推理模型

(base) PS D:\yolo_seq\TensorRT-For-YOLO-Series-cuda-python> 
python trt.py -e yolov10.trt  -i src/1.jpg -o yolov10-1.jpg --end2end

输出结果如下所示，可以发现fps此时接近200，5ms左右一张图

Namespace(end2end=True, engine='yolov10.trt', image='src/1.jpg', output='yolov10-1.jpg', video=None)
199.44339337776142 FPS

5、int8推理

导出int8模型需要设置 calib_input参数（设置校准数据的路径）

python export.py  -o yolov10n.onnx -e yolov10.int8.trt --end2end --v10 -p int8 --calib_input  D:\yolo_seq\datasets\coco128\images\train2017

然后再基于以下命令测试int8模型推理效果

python trt.py -e yolov10.int8.trt  -i src/1.jpg -o yolov10-1.jpg --end2end

此时运行输出如下

Namespace(end2end=True, engine='yolov10.int8.trt', image='src/1.jpg', output='yolov10-1.jpg', video=None)
259.70545246403356 FPS

运行效果如下，相比于fp16，似乎没有区别。

上点难度看一下int8与fp16推理效果

python trt.py -e yolov10.int8.trt  -i D:\yolo_seq\coco\images\val2017\000000001000.jpg -o 0-int8.jpg --end2end

python trt.py -e yolov10.trt  -i D:\yolo_seq\coco\images\val2017\000000001000.jpg -o 0-fp16.jpg --end2end

此时可以发现，转为int8的模型明显存在较多漏检

基于更多的数据矫正，再对比一下模型的效果差异

python export.py  -o yolov10n.onnx -e yolov10.int8-2.trt --end2end --v10 -p int8 --calib_input  D:\yolo_seq\coco\images\val2017

测试新数据，可以发现int8模型任然存在不足

python trt.py -e yolov10.int8.trt  -i D:\yolo_seq\coco\images\train2017\000000581831.jpg -o 1-int8.jpg --end2end

python trt.py -e yolov10.trt  -i D:\yolo_seq\coco\images\train2017\000000581831.jpg -o 1-fp16.jpg --end2end

测试校准数据中的样例，依旧可以发现int8模型存在精度损失

python trt.py -e yolov10.int8.trt  -i D:\yolo_seq\coco\images\val2017\000000001000.jpg -o 2-int8.jpg --end2end

python trt.py -e yolov10.trt  -i D:\yolo_seq\coco\images\val2017\000000001000.jpg -o 2-fp16.jpg --end2end