当前位置：首页 > news >正文

YOLOv5（四）：models/yolo.py

news 2025/11/11 9:52:00

一、代码逻辑架构与调用关系

1. 核心架构总览

YOLOv5的yolo.py实现了目标检测模型的完整生命周期（构建、训练、推理），核心架构分为3个主要类和1个解析函数，整体逻辑如下：
配置文件(.yaml) → parse_model解析 → Model类构建网络 → Detect类生成检测结果

2. 核心类与函数的作用及关系

组件	作用	核心依赖
`Detect`类	检测头模块，将特征图转换为边界框和类别概率	依赖输入特征图的通道数、锚点配置，输出检测结果
`Model`类	主模型类，整合backbone、neck和head，实现训练/推理流程	依赖`parse_model`生成网络层，依赖`Detect`作为检测头
`parse_model`函数	解析.yaml配置文件，生成网络层列表和需保存的层索引	依赖配置文件中的backbone/head定义，调用`models.common`中的基础层（如Conv）

3. 详细逻辑流程

（1）模型初始化流程（`Model.init`）

加载配置：加载.yaml配置文件（或字典），解析类别数nc、锚点anchors、深度因子gd、宽度因子gw等参数。
构建网络：调用parse_model函数，根据配置文件的backbone和head定义，生成网络层列表self.model和需保存的层索引self.save。
初始化检测头：计算检测层步长stride，归一化锚点，初始化检测头偏置（_initialize_biases）。
权重初始化：调用initialize_weights初始化所有层的权重。

（2）前向传播流程（`Model.forward`）

单尺度推理/训练（_forward_once）：
遍历self.model中的每一层，按层索引m.i和输入来源m.f获取输入，执行层计算并保存输出（仅self.save中的层）。
- 训练模式：返回检测头的原始特征图（用于损失计算）。
- 推理模式：检测头（Detect）将特征图转换为边界框坐标（结合网格和锚点），返回拼接后的检测结果。
增强推理（_forward_augment）：
对输入进行多尺度缩放和翻转，生成多个版本的输入，分别推理后将结果逆处理（缩放+逆翻转）并拼接，提升高检测鲁棒性。

（3）检测头解析与网络构建（`parse_model`）

参数解析：从配置文件中提取每层的from（输入来源）、n（重复次数）、m（模块类型）、args（模块参数）。
通道计算：根据宽度因子gw调整输出通道（确保能被8整除），根据深度因子gd调整重复次数n`。
模块实例化：根据模块类型（如Conv、C3、Detect）实例化参数，创建网络层并记录参数量、输入来源等信息。
输出保存：收集需要保存输出的层索引（用于后续层作为输入），最终构建nn.Sequential返回。

（4）检测头逻辑（`Detect`）

初始化：定义输出卷积层，注册锚点为缓冲区，初始化网格和锚点网格。
前向传播：
- 训练时：返回卷积后的特征图（用于损失计算）。
- 推理时：通过sigmoid激活输出，结合网格和锚点网格将特征图上的预测转换为图像坐标下的边界框（xywh）、置信度和类别概率。

4. 关键调用关系图

Model.__init__  
├─ 加载配置文件 → 解析nc、anchors等参数  
├─ parse_model() → 生成网络层self.model和保存列表self.save  
│  ├─ 遍历backbone/head配置，实例化Conv、C3等模块  
│  └─ 计算各层通道数和参数量  
├─ 初始化检测头（Detect）  
│  ├─ 计算stride（步长）  
│  └─ _initialize_biases() → 初始化检测头偏置  
└─ 初始化权重  Model.forward  
├─ 若augment=True → 调用_forward_augment（多尺度增强）  
│  └─ 多次调用_forward_once，处理后拼接结果  
└─ 否则 → 调用_forward_once_once（单尺度传播）  ├─ 遍历self.model中的层，按层计算  └─ 最终调用Detect.forward →生成检测结果  Detect.forward  
├─ 训练模式：返回特征图  
└─ 推理模式：生成网格和锚点网格，转换预测为边界框 → 返回检测结果