yolov5模型pt转wk

将 YOLOv5 模型从 PyTorch 转换为 WK(x) 格式的步骤如下：

1. 安装 ONNX：使用 pip install onnx 安装 ONNX。

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

import torch
from models.experimental import attempt_load
from utils.activations import Hardswish
import onnx
from onnxsim import simplify

# 加载 PyTorch 模型
model = attempt_load('yolov5s.pt', map_location=torch.device('cpu'))
model.eval()

# 创建模拟输入
input_shape = (1, 3, 640, 640)
inputs = torch.randn(input_shape)

# 导出 ONNX 模型
dynamic_axes = {"images": {0: "batch_size"}, "output": {0: "batch_size"}}
onnx_model_path = "yolov5s.onnx"
torch.onnx.export(model, inputs, onnx_model_path, opset_version=11, dynamic_axes=dynamic_axes, input_names=["images"], output_names=["output"], custom_opsets={"com.microsoft": 1})

# 精简 ONNX 模型
model = onnx.load(onnx_model_path)
model_simp, check = simplify(model, input_shapes={"images": input_shape}, dynamic_input_shape=True)
onnx.save(model_simp, onnx_model_path)

3. 使用 ONNX 转换器将模型转换为 WK(x) 格式：使用以下代码将 ONNX 模型转换为 WK(x) 格式：

import onnx
import onnx_coreml

# 加载 ONNX 模型
model = onnx.load('yolov5s.onnx')

# 将模型转换为 WK(x) 格式
model = onnx_coreml.convert(model, minimum_ios_deployment_target='13')
model.save('yolov5s.mlmodel')

这样就可以将 YOLOv5 模型从 PyTorch 转换为 WK(x) 格式了。