项目地址：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 YOLOmodel = 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 argparseclass Predictor(BaseEngine):def __init__(self, engine_path):super(Predictor, self).__init__(engine_path)self.n_classes = 80 # your model classesif __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.imagevideo = args.videoif 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 webcam4、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 --end2endpython 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 --end2endpython 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 --end2endpython trt.py -e yolov10.trt -i D:\yolo_seq\coco\images\val2017\000000001000.jpg -o 2-fp16.jpg --end2end
