1. 环境搭建与ONNX模型导出

任务目标

本案例完成两件事：搭建YOLOv5模型导出所需的Python虚拟环境并安装依赖；加载YOLOv5预训练权重，将其导出为ONNX格式文件，为后续部署（如ONNX Runtime推理）做好准备。

环境与依赖

• 操作系统：Windows 10/11 或 Ubuntu 20.04+
• Python版本：3.8～3.11（推荐3.9）
• PyTorch：2.0.0+（CPU或CUDA版均可）
• ONNX：1.13.0+
• onnxruntime：1.14.0+（仅用于后续验证，本案例导出不强制）

实现步骤

首先创建并激活虚拟环境，然后安装依赖。YOLOv5官方仓库的 requirements.txt 包含导出所需的全部依赖，直接使用即可。导出脚本的核心逻辑是：加载预训练权重 → 构造模型并设为eval模式 → 用随机输入执行一次前向传播以触发ONNX导出。

关键完整代码

# export_yolov5_onnx.py
import torch
import onnx
from pathlib import Path

# 1. 克隆YOLOv5仓库并安装依赖（在终端执行）
# git clone https://github.com/ultralytics/yolov5
# cd yolov5
# pip install -r requirements.txt

# 2. 导出ONNX
def export_onnx(weights='yolov5s.pt', imgsz=640, dynamic=False):
    # 加载模型（自动下载预训练权重到当前目录）
    model = torch.hub.load('ultralytics/yolov5', 'custom', path=weights, force_reload=True)
    model.eval()
    
    # 构造示例输入 [batch, 3, height, width]
    dummy_input = torch.zeros(1, 3, imgsz, imgsz)
    
    # 导出为ONNX
    export_path = Path(weights).with_suffix('.onnx')
    torch.onnx.export(
        model.model,                # 底层nn.Module
        dummy_input,
        export_path,
        opset_version=12,
        input_names=['images'],
        output_names=['output'],
        dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},
                      'output': {0: 'batch'}} if dynamic else None
    )
    
    # 验证ONNX模型完整性
    onnx_model = onnx.load(export_path)
    onnx.checker.check_model(onnx_model)
    print(f'导出成功: {export_path}')
    print(f'输入节点: {onnx_model.graph.input[0].name}, 形状: {onnx_model.graph.input[0].type.tensor_type.shape}')
    return str(export_path)

if __name__ == '__main__':
    export_onnx(weights='yolov5s.pt', imgsz=640, dynamic=False)

运行方式

# 在终端中依次执行
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt
python export_yolov5_onnx.py

首次运行会自动下载 yolov5s.pt（约14MB）到当前目录。若网络受限，可手动下载权重文件后放到脚本同目录。

验证方法与通过条件

• 文件生成：当前目录出现 yolov5s.onnx 文件，大小约28MB（动态导出约30MB）。
• 完整性校验：脚本执行后无异常，控制台输出 导出成功: yolov5s.onnx，且 onnx.checker.check_model 未报错。
• 输入形状确认：打印的输入节点名称为 images，形状维度包含 batch、3、height、width。

若需进一步验证ONNX可推理，可额外安装 onnxruntime 并执行：

import onnxruntime as ort
sess = ort.InferenceSession('yolov5s.onnx')
print(f'输入: {sess.get_inputs()[0].name}, 输出: {sess.get_outputs()[0].name}')

能成功创建会话即说明ONNX文件可被标准推理引擎加载。