如何在yolov5模型中导出onnx格式

要将 YOLOv5 模型导出为 ONNX 格式，您可以按照以下步骤进行操作：

1.安装相关的依赖库

pip install onnx onnx-simplifier

2.使用以下代码将 PyTorch 模型加载到内存中

import torch

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True, force_reload=True)

# 或者通过加载本地模型文件的方式
# model = torch.load("path/to/your/yolov5s.pt")

3.创建一个样本输入，用于模型推理

sample_input = torch.randn((1, 3, 640, 640))

4.使用以下代码将 PyTorch 模型转换为 ONNX 格式

import onnx
import onnxruntime

# 将 PyTorch 模型转换为 ONNX 模型
onnx_model = onnx.load(model.export_onnx())
onnx.checker.check_model(onnx_model)

# 简化 ONNX 模型
simplified_model = onnxsim.simplify(onnx_model)

# 将简化后的 ONNX 模型保存到文件中
onnx.save(simplified_model, 'yolov5s.onnx')

5.验证 ONNX 模型是否可用

# 创建 ONNX 运行时
ort_session = onnxruntime.InferenceSession('yolov5s.onnx')

# 使用 ONNX 模型进行推理
ort_inputs = {ort_session.get_inputs()[0].name: sample_input.cpu().numpy()}
ort_outputs = ort_session.run(None, ort_inputs)

print(ort_outputs)

其中，`model.export_onnx()` 方法用于将 PyTorch 模型导出为 ONNX 模型。在导出 ONNX 模型之后，使用 ONNX 运行时验证模型是否有效。