- 通过yolov8官方代码训练生成的权重文件，均可通过此项目转换为tensorrt量化模型，同时支持常规模型以及端到端模型
- 进一步的，同样支持输出维度结构为
# 批大小 预测框数量 x, y, w, h, obj_conf, label_scores
outputs.shape = [batch_size, number_of_preds, 5 + number_of_classes]
outputs.shape = [batch_size, 5 + number_of_classes, number_of_preds]
# 或者 x, y, w, h, scores
outputs.shape = [batch_size, number_of_preds, 4 + number_of_classes]
outputs.shape = [batch_size, 4 + number_of_classes, number_of_preds]的其他目标检测模型的onnx文件转换为tensorrt的engine文件和便于python使用的pt文件.
- CUDA, CuDNN, TensorRT（Toolkit）
- 所有requirements.txt中的三方库
注意
- end2end模型要求TensorRT版本为8.0.0及以上（本人使用的是TensorRT-8.0.1.6.Windows10.x86_64.cuda-11.3.cudnn8.2）
- 由于修改了原项目ultralytics中的一些函数，所以导出模型时不要将原项目代替本项目中的ultralytics文件夹
其中 torch2trt 安装如下
git clone https://github.com/NVIDIA-AI-IOT/torch2trt.git
cd torch2trt
python setup.py installpython 的 tensorrt 通过 TensorRT Toolkit中提供的安装包安装，然后
git clone https://github.com/LSH9832/yolov8_trt
cd yolov8_trt
pip install -r requirements.txt- 支持转onnx、tensorrt engine文件以及将engine文件与相关信息打包的，便于python中使用的pt文件
- 支持常规模型以及端到端（End2End）模型
通过yolov8官方方式下载模型
python download_all_models.py# mode有以下5个
# 1. onnx：仅导出onnx常规模型
# 2. trt：导出onnx及tensorrt常规模型
# 3. end2end：导出onnx常规模型，并通过其导出end2end模型，再转化为end2end的engine模型
# 4. onnx2trt: 其他非yolov8模型的onnx文件（未添加end2end结构）转换为tensorrt常规模型
# 5. onnx2end2end: 其他非yolov8模型的onnx文件（未添加end2end结构）转换为端到端onnx模型，再转化为end2end的engine模型
python export.py --mode trt
# 如果是yolov8模型（trt模式或end2end模式），输入以下4个参数
--weights yolov8s.pt
--batch 1
--img-size 640 640 # 4:3视频流设为 480 640， 16:9视频流设为 384 640
--opset 11
# 否则（onnx2trt模式或onnx2end2end模式）输入以下2个参数
--onnx yolov7.onnx # 其他模型的onnx文件
--cfg yolov7.yaml # 所需包含的信息见下文
# 以下所有参数如为onnx模式不用写
--workspace 10 # 转换为tensorrt模型时最大使用的显存空间(GB)
--fp16 # int8, best，模型精度使能
# 以下为end2end模式和onnx2end2end模式所需参数(Batched_NMS中所需参数)，其他模式不用填写
--conf-thres 0.2 # 置信度阈值
--nms-thres 0.6 # NMS 目标框重合度阈值
--topk 2000 # 选置信度最高的前K个目标框作为NMS的输入
--keep 100 # NMS之后最多保留前K个置信度最高的结果其中，--cfg 指向的yaml文件应该包含如下信息
batch_size: 1 # int 批大小
pixel_range: 1 # int 输入图像像素值范围， 1代表[0, 1]之间（YOLOv5,6,7等），255代表[0, 255]之间（YOLOX等）
obj_conf_enabled: true # bool 是否使用了前景预测（obj_conf）, 注意true和false开头字母用小写
img_size: [640, 640] # List[int] 输入图像 [高，宽]，注意不要反了
input_name: ["input_0"] # List[str] onnx文件输入变量名
output_name: ["output_0"] # List[str] onnx文件输出变量名
names: ["person", "car", ...] # List[str] 类别名称
不同模式导出不同模型，下面以yolov8s.pt为例
- onnx模式
./yolo_export/yolov8s/yolov8s.onnx # 常规模型
./yolo_export/yolov8s/yolov8s.yaml # 模型相关信息
./yolo_export/yolov8s/yolov8s.json # 模型相关信息- trt模式
./yolo_export/yolov8s/yolov8s.onnx # 常规模型
./yolo_export/yolov8s/yolov8s.yaml # 常规模型相关信息
./yolo_export/yolov8s/yolov8s.json # 常规模型相关信息
./yolo_export/yolov8s/yolov8s.engine # 用于C++部署
./yolo_export/yolov8s/yolov8s.pt # 用于python部署- end2end模式
./yolo_export/yolov8s/yolov8s.onnx # 常规模型
./yolo_export/yolov8s/yolov8s.yaml # 常规模型相关信息
./yolo_export/yolov8s/yolov8s.json # 常规模型相关信息
./yolo_export/yolov8s/yolov8s_end2end.onnx # 端到端模型
./yolo_export/yolov8s/yolov8s_end2end.json # 端到端模型相关信息
./yolo_export/yolov8s/yolov8s_end2end.engine # 用于C++部署
./yolo_export/yolov8s/yolov8s_end2end.pt # 用于python部署端到端的模型最后会增加如下结构, 相关代码改自https://github.com/DataXujing/YOLOv8
由本项目生成的，不论是常规模型还是端到端模型，均可使用，不依赖原项目的函数功能
python trt_infer.py --weight ./yolo_export/yolov8s/yolov8s.pt
--source path/to/your/video/file or (rtsp/rtmp/http)://xxx.xxx.xxx.xxx/live/xxx or 0
--no-label # 声明此项则不会显示类别与置信度标签，仅显示目标框
# 以下为常规模型所需参数，端到端模型无效
--conf 0.2 # 置信度阈值
--nms 0.6 # NMS 目标框重合度阈值即将推出，着急用可到其他项目中找找
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent