1. 环境搭建与模型转换基础

任务目标

本案例完成YOLOv5目标检测模型从PyTorch到ONNX的转换基础环境搭建，最终验证导出模型输出张量形状是否符合预期。核心交付物为：可运行的PyTorch环境、YOLOv5模型定义、ONNX导出脚本及形状验证逻辑。

环境依赖

  • Python 3.8+，PyTorch 1.12+（CPU或CUDA版本均可）
  • 安装依赖：pip install torch onnx onnxruntime
  • 无需完整YOLOv5仓库，本案例使用最小模型定义完成导出验证

最小模型定义与导出

以下代码定义了一个与YOLOv5检测头输出结构一致的简化模型（输入640×640，输出3个尺度的预测张量），并完成ONNX导出与形状验证。该模型保留了YOLOv5输出张量的核心维度特征：[batch, 3*(5+num_classes), grid_h, grid_w]。

import torch
import torch.nn as nn
import onnx
import onnxruntime as ort

class ConvBlock(nn.Module):
    def __init__(self, c_in, c_out, k=3, s=1):
        super().__init__()
        self.conv = nn.Conv2d(c_in, c_out, k, s, k//2, bias=False)
        self.bn = nn.BatchNorm2d(c_out)
        self.act = nn.SiLU()
    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

class YOLOv5Mini(nn.Module):
    """最小YOLOv5结构：骨干+检测头，输出3个尺度特征图"""
    def __init__(self, nc=80):
        super().__init__()
        self.backbone = nn.Sequential(
            ConvBlock(3, 32, s=2),   # 320×320
            ConvBlock(32, 64, s=2),  # 160×160
            ConvBlock(64, 128, s=2), # 80×80
            ConvBlock(128, 256, s=2),# 40×40
            ConvBlock(256, 512, s=2) # 20×20
        )
        # 检测头：3个尺度，每个anchor预测5+nc个值
        self.detect = nn.ModuleList([
            nn.Conv2d(512, 3*(5+nc), 1),  # P3: 80×80
            nn.Conv2d(256, 3*(5+nc), 1),  # P4: 40×40
            nn.Conv2d(128, 3*(5+nc), 1)   # P5: 20×20
        ])
        # 上采样层用于特征融合（简化版）
        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')

    def forward(self, x):
        # 骨干网络逐层提取特征
        c1 = self.backbone[0](x)   # 320
        c2 = self.backbone[1](c1)  # 160
        c3 = self.backbone[2](c2)  # 80
        c4 = self.backbone[3](c3)  # 40
        c5 = self.backbone[4](c4)  # 20

        # 检测头输出（简化：直接使用各层特征）
        out = []
        out.append(self.detect[0](c3))  # 80×80
        out.append(self.detect[1](c4))  # 40×40
        out.append(self.detect[2](c5))  # 20×20
        return out

# 初始化模型
model = YOLOv5Mini(nc=80)
model.eval()

# 构造输入
dummy_input = torch.randn(1, 3, 640, 640)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "yolov5_mini.onnx",
    opset_version=12,
    input_names=["images"],
    output_names=["out_80", "out_40", "out_20"],
    dynamic_axes={"images": {0: "batch"}}
)

# 验证ONNX模型完整性
onnx_model = onnx.load("yolov5_mini.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过")

# 使用onnxruntime验证输出形状
ort_session = ort.InferenceSession("yolov5_mini.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_outputs = ort_session.run(None, ort_inputs)

# 预期输出形状验证
expected_shapes = [(1, 255, 80, 80), (1, 255, 40, 40), (1, 255, 20, 20)]
for i, (out, exp_shape) in enumerate(zip(ort_outputs, expected_shapes)):
    assert out.shape == exp_shape, f"输出{i}形状{out.shape}不匹配预期{exp_shape}"
    print(f"输出{i}形状验证通过: {out.shape}")

运行与验证

  1. 将上述代码保存为 export_yolov5.py，在项目目录下执行 python export_yolov5.py
  2. 程序依次完成：模型定义→ONNX导出→结构检查→推理形状验证
  3. 通过条件：控制台输出以下内容即判定成功：
    • ONNX模型结构检查通过
    • 输出0形状验证通过: (1, 255, 80, 80)
    • 输出1形状验证通过: (1, 255, 40, 40)
    • 输出2形状验证通过: (1, 255, 20, 20)

关键说明

  • 输出通道数255 = 3×(5+80)，对应3个anchor、5个框属性(x,y,w,h,conf)和80个类别
  • 本案例使用简化模型验证导出链路，实际YOLOv5s模型结构更复杂，但导出流程与验证方法完全一致
  • 若需导出完整YOLOv5s，可替换模型定义部分，其余导出与验证代码无需修改
← 上一章 📖 返回目录 下一章 →