使用TensorRT加速yolov5

版本:

TensorRT:8.2.5.1

1.安装TensorRT,上一篇博客有教程https://blog.csdn.net/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如图:
使用TensorRT加速yolov5_第1张图片

解决方法:

打开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文件。

你可能感兴趣的:(深度学习,python,目标检测)