版本:
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
如图:
解决方法:
打开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
文件。