### 将 YOLOv5 的 PyTorch (.pt) 模型转换为 ONNX 格式的详细方法 为了将 YOLOv5 的 `.pt` 文件转换为 `.onnx` 格式，以下是完整的实现流程： #### 安装必要的库 确保安装了 `torch` 和 `onnx` 库。可以通过以下命令完成安装： ```bash pip install torch torchvision onnx ``` #### 编写转换脚本 创建一个 Python 脚本（例如 `convert_yolo_to_onnx.py`），并将以下代码复制进去： ```python import torch from torch.utils.mobile_optimizer import optimize_for_mobile # 加载自定义的 YOLOv5 模型 model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt') # 修改路径为你自己的 .pt 文件路径 model.eval() # 设置输入张量形状 (batch_size, channels, height, width) dummy_input = torch.randn(1, 3, 640, 640) # 导出模型到 ONNX 格式 output_file = 'yolov5s.onnx' torch.onnx.export( model, dummy_input, output_file, opset_version=12, # 使用支持的 OPSET 版本 input_names=['images'], # 输入节点名称 output_names=['output'] # 输出节点名称 ) print(f"Model has been successfully converted to ONNX format and saved as {output_file}.") ``` #### 运行脚本 保存并运行该脚本即可完成转换： ```bash python convert_yolo_to_onnx.py ``` 完成后，在当前工作目录下会生成一个名为 `yolov5s.onnx` 的文件。 --- #### 验证转换结果 验证 ONNX 模型是否成功导出，可以加载模型并检查其结构： ```python import onnx # 加载 ONNX 模型 onnx_model = onnx.load("yolov5s.onnx") # 检查模型结构 onnx.checker.check_model(onnx_model) print("ONNX Model is valid.") ``` 如果输出显示 `"ONNX Model is valid."`，说明转换成功[^3]。 --- #### 注意事项 1. **输入尺寸匹配**：YOLOv5 默认输入图像大小为 `(640, 640)`，因此在设置 `dummy_input` 时需保持一致。 2. **OPSET 版本兼容性**：推荐使用 `opset_version=12` 或更高版本以获得更好的兼容性。 3. **依赖环境一致性**：确保使用的 PyTorch 版本与 ONNX 支持的版本相匹配[^1]。 ---