Yolov5使用TensorRT加速(python篇)/win10/linux

本文介绍如何在最新项目中利用TensorRT简化权重文件导出和模型加速过程,包括安装TensorRT、更新Yolov5版本、导出.engine文件以及使用步骤,重点在环境配置和关键操作上.

前言

使用的是最新的项目,因为新的项目都把代码写好了,直接调用就行了。。。意思操作更无脑了。
前提需要安装TensorRT,使用操作只需要两步。。。麻烦在于环境的安装。

加速流程:导出.engine格式权重文件,加载engine文件自动调用tensorRT。

〇、准备工作:

1.安装TensorRT

参考我另一篇博客超简单安装TensorRT(仅限python)

2.更新项目

使用的最新版的yolov5 6.0或6.1
没更新的更新项目:
git clone https://ghproxy.com/https://github.com/ultralytics/yolov5

3.检查依赖

一定要记得激活环境!
conda actiavate <环境名>
pip install -r requirements.txt
不存在的包会自动更新。

二、导出权重文件

.pt导出为.engine
形式:
python export.py --weights pt权重文件路径 --data 你数据集yaml文件路径 --include engine --device 0

  • weights 权重文件路径
  • data 数据集yaml文件路径
  • include 要导出的模型
  • device 使用设备,0代表使用gpu

举例:
python export.py --weights models/weights/demo.pt --data data/data.yaml --include engine --device 0 --half

-- half 开启半精度推理会快一些

如果导出过程中报错,说哪个库没有装,就去装完库在执行上面导出命令

导出的结果和自己的pt文件在一个路径下。

三、使用tensorRT加速

python detect.py --source xxx/xxx/valid/images --data xxx/xxxx/xx.yaml --weights xxx/xxx/demo.engine --device 0

### 使用 TensorRT 加速 YOLOv5 v6.0 部署和推理 #### 准备工作 为了确保能够顺利部署并利用 TensorRTYOLOv5 进行加速,需先完成必要的前置准备工作。这包括但不限于安装 Python 虚拟环境、设置 OpenCV 和安装 TensorRT。 对于 Windows 用户而言,在 Win10 系统环境下,建议创建一个新的 Anaconda 或者其他形式的 Python 虚拟环境来管理依赖项[^1]。而对于 Linux (Ubuntu_x86_64),同样推荐建立独立的 Python 环境用于项目开发,并通过 pip 工具安装所需的库文件如 NumPy, PyTorch 及其对应的 CUDA 版本等[^2]。 #### 安装 TensorRT 无论是哪种操作系统平台,都需要下载适合当前系统的 TensorRT SDK 并按照官方文档指导完成安装流程。特别注意要匹配好所使用的 GPU 架构以及 CUDA/CUDNN 的版本号以避免兼容性问题的发生[^3]。 #### 导出 ONNX 模型 YOLOv5 默认训练得到的是 .pt 后缀名保存下来的模型权重文件。为了让 TensorRT 更高效地加载这些预训练好的网络结构参数,通常会先把它们转换成通用中间表示格式——ONNX 文件(.onnx)[^1]。此操作可以通过调用 `torch.onnx.export()` 方法实现: ```python import torch from models.experimental import attempt_load device = 'cuda' if torch.cuda.is_available() else 'cpu' model = attempt_load('path/to/yolov5.pt', map_location=device) # Load model dummy_input = torch.zeros(1, 3, 640, 640).to(device) torch.onnx.export( model, dummy_input, "yolov5.onnx", opset_version=11, do_constant_folding=True, input_names=['input'], output_names=['output'] ) ``` 这段脚本将会把指定路径下的 `.pt` 格式的 YOLOv5 权重导出为名为 `yolov5.onnx` 的 ONNX 文件。 #### 编译 TRT Engine 有了上述生成的 ONNX 文件之后就可以借助于 TensorRT 提供的相关工具将其编译成为专有的运行时引擎(TRTEngine),从而达到优化性能的目的。具体来说可以采用 trtexec 命令行工具来进行这项任务: ```bash trtexec --onnx=yolov5.onnx --saveEngine=yolov5.trt --workspace=1024 --fp16 ``` 这条命令将基于给定的 ONNX 描述构建一个 FP16 半精度浮点数运算模式下的 TensorRT 引擎,并存储到本地磁盘作为后续执行的基础。 #### 推理测试 最后一步就是编写简单的 C++/Python 应用来读取之前制作出来的 TensorRT 引擎并对输入数据实施前向传播计算获取检测结果。这里给出一段简化的 Python 示例代码片段展示如何使用 pycuda 和 tensorrt python api 结合的方式快速启动图像识别服务: ```python import numpy as np import cv2 import pycuda.autoinit import pycuda.driver as cuda import tensorrt as trt def load_engine(trt_runtime, plan_path): with open(plan_path, 'rb') as f: engine_data = f.read() engine = trt_runtime.deserialize_cuda_engine(engine_data) return engine if __name__ == '__main__': logger = trt.Logger(trt.Logger.WARNING) runtime = trt.Runtime(logger) engine_file_path = "./yolov5.trt" # Load the serialized ICudaEngine from file and create an execution context. engine = load_engine(runtime, engine_file_path) context = engine.create_execution_context() # Allocate buffers to hold the I/O of the network. inputs, outputs, bindings, stream = allocate_buffers(engine) image = cv2.imread("test.jpg") # Replace this line according your actual situation. # Preprocess the data before feeding into inference API... processed_image = preprocess(image) # Perform Inference using optimized TensorRT engine here... detections = post_process(outputs[0].host) draw_boxes_on_image(detections, image) ``` 以上即完成了整个从准备阶段直至最终应用层面的一系列步骤介绍,帮助理解怎样有效地运用 TensorRT 技术提升 YOLOv5 在目标检测场景中的表现效率。
评论 40
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值