【源码】基于TensorRT的YOLOv5目标检测

项目简介

本项目旨在借助TensorRT对YOLOv5目标检测模型进行部署和优化，利用GPU资源提高目标检测的速度与精度，实现高效的目标检测任务。

项目的主要特性和功能

1. 具备模型导出与加载功能，可将YOLOv5模型导出为ONNX格式，并通过TensorRT加载ONNX模型进行推断。
2. 支持动态形状输入，能在运行时动态调整输入图像尺寸，适应不同分辨率的输入。
3. 集成了TensorRT的BatchNMS插件，用于非极大值抑制，减少手动编写NMS代码的工作量。
4. 支持FP16和INT8优化，在提高模型推断速度的同时保持一定的精度。

安装使用步骤

环境准备

• 系统：推荐使用Ubuntu20.04LTS或TensorRT官方Docker镜像。
• 软件版本：TensorRT 7.2.3、截止2021.09.11的Yolov5 develop分支代码、gcc 9.3.0、torch 1.8.2、onnx 1.10.1、onnx - simplifier 0.3.6。

模型训练

复制Yolov5的官方代码，按照教程训练得到pt权重文件。

模型导出

使用models/export.py脚本将YOLOv5的pt模型转换为ONNX模型，示例命令如下： bash python models/export.py --weights yolov5s.pt --img-size 640 --batch-size 1 --device 0 --include onnx --inplace --dynamic --simplify --opset-version 11 --img test_img/1.jpg

编译C++代码

复制yolov5_trt代码到本地，进入项目目录并编译： bash cd yolov5_trt mkdir build cd build cmake .. make -j 10

配置YAML

根据项目需求配置YAML文件，指定模型、输入和输出文件路径等，示例配置如下： yaml path: "../data/person/" model: onnx: "person.onnx" tensorrt: "person.trt" image: input: "person.jpg" output: "person_result.jpg" args: n_classes: 1 names: "person.names" channels: 3 height: 640 width: 640 ifdynamic: 1 min: [1, 3, 640, 640] opt: [1, 3, 640, 640] max: [1, 3, 640, 640] build: demo: 1 fp16: 1 int8: 0 workspace: 4 nms: topK: 512 keepTopK: 100 clipBoxes: 0 iouThreshold: 0.25 scoreThreshold: 0.45 isNormalized: false output: ["num_detections", "nmsed_boxes", "nmsed_scores", "nmsed_classes"]

运行程序

进入build文件夹，执行生成的可执行文件，并指定YAML配置文件的路径： bash ./yolov5_trt -c ../data/person/person.yaml

注意事项

• 根据实际环境调整环境准备步骤中的软件版本。
• 根据项目需求调整YAML配置文件中的参数。
• 在模型导出时，根据模型的结构和输入尺寸，调整models/export.py脚本中的参数。
• 使用TensorRT进行模型优化时，注意平衡模型的精度和性能。
• 根据实际情况调整模型推断的批次大小，以获得最佳的推断性能。

下载地址

点击下载 【提取码: 4003】