使用TensorRT加速yolov5

最新推荐文章于 2024-01-30 10:30:11 发布

氧气邮递员001

最新推荐文章于 2024-01-30 10:30:11 发布

阅读量1k

点赞数 1

CC 4.0 BY-SA版权

文章标签：深度学习 python 目标检测

本文链接：https://blog.youkuaiyun.com/Power392413/article/details/126321707

本文介绍如何安装TensorRT 8.2.5.1，并使用YOLOv5导出engine模型。针对导出过程中可能出现的silu操作不支持的问题提供了具体的解决方案。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

版本：

TensorRT：8.2.5.1

1.安装TensorRT，上一篇博客有教程https://blog.youkuaiyun.com/Power392413/article/details/126320255

2.下载yolov5

git clone https://github.com/ultralytics/yolov5

3.导出engine模型

conda activate ddaction
cd yolov5
python export.py --weights yolov5s.pt --device 0 --include engine

`如果出现如下错误：`

RuntimeError: Exporting the operator silu to ONNX opset version 11 is not supported. Please open a bug to request ONNX export support for the missing operator如图：
在这里插入图片描述

解决方法：

打开anaconda3\envs\ddaction\Lib\site-packages\torch\nn\modules\activation.py文件,找到class SiLU(Module):类中的def forward(self, input: Tensor) -> Tensor:方法进行修改：

def forward(self, input: Tensor) -> Tensor:
	# return F.silu(input, inplace=self.inplace)
    return input * torch.sigmoid(input)

修改完成后重新执行：

python export.py --weights yolov5s.pt --device 0 --include engine

执行完成后会在当前文件夹生成yolov5s.onnx文件和yolov5s.engine文件。