1. 项目概述
最近YOLOv6正式发布,作为目标检测领域的新成员,它在精度和速度上都有显著提升。作为一名长期从事计算机视觉开发的工程师,我决定在Windows平台上尝试部署YOLOv6的TensorRT模型,并成功实现了图片和视频的推理功能。整个过程涉及模型转换、环境配置、代码编译等多个环节,虽然踩了不少坑,但最终效果令人满意。
这个项目适合有一定C++和深度学习基础的开发者参考。通过本文,你将学会如何将YOLOv6模型从PyTorch格式转换为TensorRT引擎,并使用C++编写高效的推理程序。相比Python实现,C++版本在性能上更有优势,特别适合需要低延迟的应用场景。
2. 环境准备
2.1 硬件要求
首先确认你的电脑满足以下硬件条件:
- Windows 10或11操作系统
- NVIDIA独立显卡(非集成显卡)
- 至少8GB显存(RTX2070及以上推荐)
注意:可以通过任务管理器(Ctrl+Shift+Esc)查看显卡信息,必须看到"NVIDIA"字样才表示有独立显卡。AMD显卡或集成显卡无法使用CUDA加速。
2.2 软件环境
以下是经过测试的软件版本组合:
- Visual Studio 2019/2022
- CMake 3.30.1
- CUDA 11.8.0 + cuDNN 8.9.7
- TensorRT 8.6.1.6
- OpenCV 4.9.0
- Anaconda3 + Python 3.10
- PyTorch 2.5.1+cu124
建议尽量保持版本一致以避免兼容性问题。安装这些组件时需要注意:
- CUDA和cuDNN版本必须严格匹配
- TensorRT版本要与CUDA版本兼容
- OpenCV建议使用预编译版本,配置环境变量更方便
3. 模型转换流程
3.1 获取YOLOv6模型
首先从官方仓库下载预训练的YOLOv6模型(如yolov6n.pt)。这个PyTorch模型文件是我们转换的起点。
3.2 转换为ONNX格式
在Anaconda环境中安装好YOLOv6所需依赖后,执行以下转换命令:
bash复制yolo export model=yolov6n.pt format=onnx dynamic=False opset=12
关键参数说明:
dynamic=False:生成静态尺寸的ONNX模型,便于后续优化opset=12:指定ONNX算子集版本,确保兼容性
转换成功后你会得到yolov6n.onnx文件。
3.3 编译C++推理代码
项目使用CMake构建,需要修改CMakeLists.txt中的路径配置:
cmake复制# Find and include OpenCV
set(OpenCV_DIR "D:\\lufiles\\opencv490\\build\\x64\\vc16\\lib")
# Set TensorRT path if not set in environment variables
set(TENSORRT_DIR "D:\\lufiles\\TensorRT-8.6.1.6")
然后执行标准CMake构建流程:
bash复制mkdir build
cd build
cmake ..
在VS中打开生成的解决方案文件,选择x64 Release配置编译项目。
4. TensorRT引擎生成
4.1 使用trtexec转换
编译完成后,使用TensorRT自带的trtexec工具将ONNX模型转换为TensorRT引擎:
bash复制trtexec --onnx=yolov6n.onnx --saveEngine=yolov6n.engine --fp16
重要参数说明:
--fp16:启用FP16精度,可显著提升推理速度- 转换过程可能需要20-30分钟,取决于模型复杂度和硬件性能
4.2 引擎特性说明
生成的.engine文件有几点需要注意:
- 硬件绑定:引擎文件与生成它的GPU硬件紧密相关,不能直接在不同型号GPU间共享
- 精度选择:FP16在保持较好精度的同时大幅提升速度,是推荐选择
- 动态形状:本例使用静态尺寸,如需动态输入需修改ONNX导出参数
5. 推理测试
5.1 图片推理
使用编译好的可执行文件测试图片:
bash复制yolov6-tensorrt.exe yolov6n.engine "test.jpg"
程序会输出检测结果并保存可视化图片。典型输出包括:
- 检测到的物体类别和置信度
- 边界框坐标
- 推理耗时(包括预处理、推理、后处理)
5.2 视频推理
视频测试命令类似:
bash复制yolov6-tensorrt.exe yolov6n.engine "car.mp4"
程序会逐帧处理视频并显示实时结果。在RTX2070上,YOLOv6n模型通常能达到100+ FPS的性能。
6. 性能优化技巧
6.1 模型层面优化
- 尝试不同的YOLOv6变体(n/s/m/l)平衡精度和速度
- 使用TensorRT的INT8量化进一步加速(需校准数据集)
- 调整输入分辨率,较小的尺寸速度更快但精度可能下降
6.2 代码层面优化
- 使用异步推理重叠计算和IO操作
- 批量处理多帧提高GPU利用率
- 优化后处理逻辑,避免不必要的内存拷贝
7. 常见问题解决
7.1 环境配置问题
问题: 编译时报错找不到CUDA/cuDNN
解决:
- 检查环境变量PATH是否包含CUDA和cuDNN的bin目录
- 确认VS项目中CUDA工具包版本正确
- 重新运行CMake并清除缓存
7.2 模型转换问题
问题: trtexec转换失败,显示不支持的算子
解决:
- 检查ONNX算子集版本是否兼容
- 尝试更新TensorRT到最新版本
- 对于不支持的算子,考虑自定义插件实现
7.3 推理性能问题
问题: 推理速度远低于预期
解决:
- 确认使用的是Release模式编译
- 检查GPU利用率,确保没有其他程序占用资源
- 尝试启用FP16或INT8精度
8. 二次开发指南
如需基于此项目进行二次开发,需要注意:
- 模型输入输出结构:了解YOLOv6的输入尺寸和输出格式
- 后处理逻辑:解码边界框和非极大值抑制的实现
- 内存管理:合理分配和释放GPU内存,避免泄漏
- 多线程处理:合理设计流水线提高吞吐量
项目源码中包含详细的注释,建议从main函数开始逐步理解整个流程。对于不熟悉C++的开发者,建议先掌握基本的CUDA和TensorRT编程概念。