以下代码位于yolov5/models/yolo.py
一、导包模块
# Ultralytics YOLOv5 🚀, AGPL-3.0 license
"""
# 这是YOLOv5项目的一部分，遵循AGPL-3.0开源许可证。
# 该文件包含YOLOv5特有的模块定义和一些实用工具函数。
# 使用说明：
# $ python models/yolo.py --cfg yolov5s.yaml
# 上述命令行用于演示如何使用本文件中的模块来加载并运行YOLOv5模型，
# 其中'yolov5s.yaml'是模型的配置文件。
# 导入必要的Python库和模块
import argparse # 用于解析命令行参数
import contextlib # 提供上下文管理器
import math # 提供数学函数
import os # 操作系统接口
import platform # 获取平台信息
import sys # 访问或修改解释器变量
from copy import deepcopy # 复制模块，用于深复制对象
from pathlib import Path # 文件系统路径操作
# 获取当前文件的绝对路径
FILE = Path(__file__).resolve()
# 定义YOLOv5的根目录
ROOT = FILE.parents[1] # YOLOv5的根目录是当前文件的上两级目录
# 将YOLOv5的根目录添加到系统路径中，以便可以从中导入其他模块
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT))
# 如果当前系统不是Windows，则将根目录设置为相对路径
if platform.system() != "Windows":
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # 相对于当前工作目录的相对路径
# 从YOLOv5的其他文件中导入各种模块和类
# 这些模块和类是YOLOv5架构中不同组件的实现
from models.common import (
C3, C3SPP, C3TR, SPP, SPPF, Bottleneck, BottleneckCSP, C3Ghost, C3x, Classify, Concat, Contract, Conv,
CrossConv, DetectMultiBackend, DWConv, DWConvTranspose2d, Expand, Focus, GhostBottleneck, GhostConv, Proto
)
# 导入YOLOv5的一些实用工具函数
from utils.autoanchor import check_anchor_order # 检查锚点顺序的正确性
from utils.general import ( # 各种通用的辅助函数
LOGGER, check_version, check_yaml, colorstr, make_divisible, print_args
)
from utils.plots import feature_visualization # 特征可视化工具
from utils.torch_utils import ( # PyTorch相关的辅助函数
fuse_conv_and_bn, initialize_weights, model_info, profile, scale_img, select_device, time_sync
)
# 尝试导入thop模块，用于计算网络的FLOPs
# 如果模块不存在，thop将被设为None
try:
import thop
except ImportError:
thop = None
二、检测头
class Detect(nn.Module):
"""
YOLOv5的检测头，用于检测模型。
"""
stride = None # 在构建时计算的步长
dynamic = False # 强制重新构造网格
export = False # 导出模式
def __init__(self, nc=80, anchors=(), ch=(), inplace=True):
"""
初始化YOLOv5检测层。
参数:
nc (int): 类别数量，默认为80。
anchors (tuple): 锚框列表。
ch (tuple): 输入通道数列表。
inplace (bool): 是否使用原地操作。
"""
super().__init__() # 调用父类nn.Module的初始化方法
self.nc = nc # 类别数量
self.no = nc + 5 # 每个锚框的输出数量
self.nl = len(anchors) # 检测层数量
self.na = len(anchors[0]) // 2 # 每个层级的锚框数量
self.grid = [torch.empty(0) for _ in range(self.nl)] # 初始化网格列表
self.anchor_grid = [torch.empty(0) for _ in range(self.nl)] # 初始化锚框网格列表
self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2)) # 注册锚框张量
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # 输出卷积层列表
self.inplace = inplace # 是否使用原地操作标志
def forward(self, x):
"""
前向传播函数，处理输入数据并生成检测结果。
参数:
x (list[Tensor]): 模型的特征图列表。
返回:
list[Tensor]: 检测结果列表，每个元素对应一个层级的输出。
"""
z = [] # 初始化用于存储检测输出的列表
for i in range(self.nl): # 遍历每个检测层级
x[i] = self.m[i](x[i]) # 卷积操作
bs, _, ny, nx = x[i].shape # 获取批次大小、通道数、高度和宽度
# 调整张量形状为(batch, anchors, grid_height, grid_width, outputs)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
if not self.training: # 如果在推理阶段
if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
# 根据Segment类或Detect类的不同，分别处理输出
if isinstance(self, Segment):
xy, wh, conf, mask = x[i].split((2, 2, self.nc + 1, self.no - self.nc - 5), 4)
xy = (xy.sigmoid() * 2 + self.grid[i]) * self.stride[i] # 解码xy坐标
wh = (wh.sigmoid() * 2) ** 2 * self.anchor_grid[i] # 解码wh尺寸
y = torch.cat((xy, wh, conf.sigmoid(), mask), 4) # 合并输出
else:
xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4)
xy = (xy * 2 + self.grid[i]) * self.stride[i] # 解码xy坐标
wh = (wh * 2) ** 2 * self.anchor_grid[i] # 解码wh尺寸
y = torch.cat((xy, wh, conf), 4) # 合并输出
# 将输出张量reshape为(batch, num_anchors * grid_h * grid_w, num_outputs)
z.append(y.view(bs, self.na * nx * ny, self.no))
# 根据训练/导出模式返回不同的格式
return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, "1.10.0")):
"""
生成网格和锚框网格，兼容不同版本的PyTorch。
参数:
nx (int): 网格宽度。
ny (int): 网格高度。
i (int): 当前检测层级索引。
torch_1_10 (bool): PyTorch版本是否大于等于1.10。
返回:
tuple[Tensor, Tensor]: 网格张量和锚框网格张量。
"""
d = self.anchors[i].device # 设备类型
t = self.anchors[i].dtype # 数据类型
shape = 1, self.na, ny, nx, 2 # 目标网格形状
# 创建网格张量
y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
yv, xv = torch.meshgrid(y, x, indexing="ij") if torch_1_10 else torch.meshgrid(y, x)
grid = torch.stack((xv, yv), 2).expand(shape) - 0.5 # 创建网格
# 创建锚框网格张量
anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
return grid, anchor_grid
三、分割头
class Segment(Detect):
"""
YOLOv5的分割头，用于分割模型。
"""
def __init__(self, nc=80, anchors=(), nm=32, npr=256, ch=(), inplace=True):
"""
初始化YOLOv5分割头。
参数:
nc (int): 类别数量，默认为80。
anchors (tuple): 锚框列表。
nm (int): 掩模数量，默认为32。
npr (int): 原型数量，默认为256。
ch (tuple): 输入通道数列表。
inplace (bool): 是否使用原地操作。
"""
# 调用父类Detect的初始化方法，继承其属性和功能
super().__init__(nc, anchors, ch, inplace)
self.nm = nm # 掩模数量
self.npr = npr # 原型数量
self.no = 5 + nc + self.nm # 每个锚框的输出数量（包含掩模）
# 更新输出卷积层以适应新的输出数量
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)
# 添加原型模块，用于生成原型掩模
self.proto = Proto(ch[0], self.npr, self.nm)
# 重定义detect属性，直接指向父类的forward方法，以便在forward中使用
self.detect = Detect.forward
def forward(self, x):
"""
前向传播函数，处理输入数据并生成检测结果和原型掩模。
参数:
x (list[Tensor]): 模型的特征图列表。
返回:
tuple: 包含检测结果和原型掩模的元组，根据训练/导出模式调整输出。
"""
# 通过原型模块生成原型掩模
p = self.proto(x[0])
# 调用父类的前向传播方法来获取检测结果
x = self.detect(self, x)
# 根据训练/导出模式调整输出
if self.training: # 训练模式下返回检测结果和原型掩模
return x, p
elif self.export: # 导出模式下仅返回检测结果和原型掩模
return x[0], p
else: # 其他模式下返回检测结果、原型掩模以及额外的输出（如果有的话）
return x[0], p, x[1]
四、基础类模型
class BaseModel(nn.Module):
"""YOLOv5的基础模型类，继承自PyTorch的nn.Module."""
def forward(self, x, profile=False, visualize=False):
"""执行YOLOv5基础模型的单尺度推理或训练过程，可选择开启性能分析和特征可视化.
参数:
x (Tensor): 输入张量.
profile (bool): 是否进行性能分析.
visualize (bool): 是否启用特征可视化.
返回:
Tensor: 模型的输出.
"""
return self._forward_once(x, profile, visualize) # 单尺度推理或训练.
def _forward_once(self, x, profile=False, visualize=False):
"""执行YOLOv5模型的一次前向传播，允许性能分析和特征可视化选项.
参数:
x (Tensor): 输入张量.
profile (bool): 是否进行性能分析.
visualize (bool): 是否启用特征可视化.
返回:
Tensor: 最终输出张量.
"""
y, dt = [], [] # 保存各层输出和时间差
for m in self.model: # 遍历模型中的每一层
if m.f != -1: # 如果层不是从上一层接收输入
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]
if profile: # 如果需要性能分析
self._profile_one_layer(m, x, dt)
x = m(x) # 执行层的操作
y.append(x if m.i in self.save else None) # 保存输出，如果需要
if visualize: # 如果需要特征可视化
feature_visualization(x, m.type, m.i, save_dir=visualize)
return x # 返回最终输出
def _profile_one_layer(self, m, x, dt):
"""对单层进行性能分析，计算GFLOPs、执行时间和参数数量.
参数:
m (Module): 当前层.
x (Tensor): 输入张量.
dt (list): 存储每层的时间差.
"""
c = m == self.model[-1] # 是否是最后一层，用于防止inplace操作
o = thop.profile(m, inputs=(x.copy() if c else x), verbose=False)[0] / 1e9 * 2 if thop else 0
t = time_sync() # 同步时间
for _ in range(10): # 运行10次以获得平均时间
m(x.copy() if c else x)
dt.append((time_sync() - t) * 100) # 计算时间差
if m == self.model[0]: # 如果是第一层，打印标题
LOGGER.info("time (ms) GFLOPs params module")
LOGGER.info(f"{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f} {m.type}") # 打印时间、GFLOPs和参数量
if c: # 如果是最后一层，打印总时间
LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s} Total")
def fuse(self):
"""融合Conv2d和BatchNorm2d层以提高推理速度.
返回:
self: 修改后的模型.
"""
LOGGER.info("Fusing layers... ")
for m in self.model.modules(): # 遍历所有模块
if isinstance(m, (Conv, DWConv)) and hasattr(m, "bn"): # 如果是Conv或DWConv且有BN层
m.conv = fuse_conv_and_bn(m.conv, m.bn) # 融合卷积和BN
delattr(m, "bn") # 删除BN属性
m.forward = m.forward_fuse # 使用融合后的前向传播
self.info() # 打印模型信息
return self # 返回自身
def info(self, verbose=False, img_size=640):
"""打印模型信息，包括详细程度和输入图像大小.
参数:
verbose (bool): 是否详细打印.
img_size (int): 图像大小.
"""
model_info(self, verbose, img_size) # 调用model_info函数
def _apply(self, fn):
"""应用变换如to(), cpu(), cuda(), half()到模型张量，但不包括参数或注册缓冲区.
参数:
fn (function): 要应用的函数.
返回:
self: 修改后的模型.
"""
self = super()._apply(fn) # 应用变换到模型
m = self.model[-1] # 获取最后一个模块，通常是检测器
if isinstance(m, (Detect, Segment)): # 如果是检测或分割模块
m.stride = fn(m.stride) # 应用变换到stride
m.grid = list(map(fn, m.grid)) # 应用变换到grid
if isinstance(m.anchor_grid, list):
m.anchor_grid = list(map(fn, m.anchor_grid)) # 应用变换到anchor_grid
return self # 返回自身
五、检测模型类
class DetectionModel(BaseModel):
# YOLOv5 detection model
def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, anchors=None):
# 构造函数初始化YOLOv5模型，接受配置文件名、输入通道数、类别数量和自定义锚点。
super().__init__() # 调用基类构造函数
if isinstance(cfg, dict): # 如果cfg是一个字典，说明是已经解析过的模型配置
self.yaml = cfg # 将字典赋值给self.yaml
else: # 否则，假设cfg是一个指向YAML配置文件的路径
import yaml # 导入YAML库用于读取配置文件
self.yaml_file = Path(cfg).name # 获取配置文件名
with open(cfg, encoding="ascii", errors="ignore") as f: # 打开并读取配置文件
self.yaml = yaml.safe_load(f) # 加载YAML配置文件到self.yaml
# 定义模型
ch = self.yaml["ch"] = self.yaml.get("ch", ch) # 输入通道数，如果配置中有则使用，否则使用默认值
if nc and nc != self.yaml["nc"]: # 如果传递了类别数量并且与配置中的不同
LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}") # 记录覆盖信息
self.yaml["nc"] = nc # 更新配置文件中的类别数量
if anchors: # 如果传递了自定义锚点
LOGGER.info(f"Overriding model.yaml anchors with anchors={anchors}") # 记录覆盖信息
self.yaml["anchors"] = round(anchors) # 更新配置文件中的锚点
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # 解析模型配置并构建模型
self.names = [str(i) for i in range(self.yaml["nc"])] # 默认类别名称列表
self.inplace = self.yaml.get("inplace", True) # 是否使用原位运算
# 构建步长和锚点
m = self.model[-1] # 获取模型的最后一个模块（通常是Detect或Segment）
if isinstance(m, (Detect, Segment)): # 如果最后一个模块是Detect或Segment
def _forward(x): # 定义一个内部函数来前向传播
return self.forward(x)[0] if isinstance(m, Segment) else self.forward(x)
s = 256 # 最小步长的两倍
m.inplace = self.inplace # 设置模块的原位运算属性
m.stride = torch.tensor([s / x.shape[-2] for x in _forward(torch.zeros(1, ch, s, s))]) # 计算步长
check_anchor_order(m) # 检查锚点顺序
m.anchors /= m.stride.view(-1, 1, 1) # 调整锚点大小
self.stride = m.stride # 设置模型的步长属性
self._initialize_biases() # 初始化偏置
# 初始化权重和偏置
initialize_weights(self) # 初始化模型权重
self.info() # 输出模型信息
LOGGER.info("") # 输出空行分隔日志
def forward(self, x, augment=False, profile=False, visualize=False):
# 执行单尺度或增强推断，可能包括性能分析或可视化。
if augment:
return self._forward_augment(x) # 增强推断
return self._forward_once(x, profile, visualize) # 单尺度推断
def _forward_augment(self, x):
# 在不同的尺度和翻转下执行增强推断，返回组合后的检测结果。
img_size = x.shape[-2:] # 图像的高度和宽度
s = [1, 0.83, 0.67] # 不同的缩放比例
f = [None, 3, None] # 翻转类型（无翻转，水平翻转，垂直翻转）
y = [] # 存储输出
for si, fi in zip(s, f): # 遍历不同的尺度和翻转
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max())) # 缩放和翻转图像
yi = self._forward_once(xi)[0] # 前向传播
yi = self._descale_pred(yi, fi, si, img_size) # 反缩放预测结果
y.append(yi) # 添加到输出列表
y = self._clip_augmented(y) # 裁剪增强推断的尾巴
return torch.cat(y, 1), None # 返回拼接后的输出
def _descale_pred(self, p, flips, scale, img_size):
# 反缩放增强推断的预测结果，调整翻转和图像尺寸。
if self.inplace: # 如果使用原位运算
p[..., :4] /= scale # 反缩放边界框坐标
if flips == 2: # 如果进行了垂直翻转
p[..., 1] = img_size[0] - p[..., 1] # 反翻转y坐标
elif flips == 3: # 如果进行了水平翻转
p[..., 0] = img_size[1] - p[..., 0] # 反翻转x坐标
else: # 如果不使用原位运算
x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale # 分离坐标和宽高
if flips == 2: # 如果进行了垂直翻转
y = img_size[0] - y # 反翻转y坐标
elif flips == 3: # 如果进行了水平翻转
x = img_size[1] - x # 反翻转x坐标
p = torch.cat((x, y, wh, p[..., 4:]), -1) # 重新组合坐标和宽高
return p # 返回反缩放后的预测结果
def _clip_augmented(self, y):
# 裁剪增强推断的尾巴，影响第一个和最后一个张量基于网格点和层数。
nl = self.model[-1].nl # 检测层数
g = sum(4**x for x in range(nl)) # 总网格点数
e = 1 # 排除层数计数
i = (y[0].shape[1] // g) * sum(4**x for x in range(e)) # 大尺度裁剪索引
y[0] = y[0][:, :-i] # 大尺度裁剪
i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x in range(e)) # 小尺度裁剪索引
y[-1] = y[-1][:, i:] # 小尺度裁剪
return y # 返回裁剪后的结果
def _initialize_biases(self, cf=None):
# 初始化YOLOv5的Detect()模块的偏置，可选地使用类别频率。
m = self.model[-1] # 获取Detect模块
for mi, s in zip(m.m, m.stride): # 遍历模块和步长
b = mi.bias.view(m.na, -1) # 查看偏置为(锚点数, 类别数+5)
b.data[:, 4] += math.log(8 / (640 / s) ** 2) # 对象偏置初始化
b.data[:, 5 : 5 + m.nc] += (
math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
) # 类别偏置初始化
mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) # 更新偏置参数
Model = DetectionModel # 保留YOLOv5 'Model'类以便向后兼容
六、分割模型类
class SegmentationModel(DetectionModel):
# YOLOv5 segmentation model
def __init__(self, cfg="yolov5s-seg.yaml", ch=3, nc=None, anchors=None):
"""Initializes a YOLOv5 segmentation model with configurable params: cfg (str) for configuration, ch (int) for channels, nc (int) for num classes, anchors (list)."""
super().__init__(cfg, ch, nc, anchors)
七、分类模型类
# 定义ClassificationModel类，用于YOLOv5分类任务，继承自BaseModel。
#
class ClassificationModel(BaseModel):
# YOLOv5 classification model
def __init__(self, cfg=None, model=None, nc=1000, cutoff=10):
"""Initializes YOLOv5 model with config file `cfg`, input channels `ch`, number of classes `nc`, and `cuttoff`
index.
"""
super().__init__() # 调用基类构造器
self._from_detection_model(model, nc, cutoff) if model is not None else self._from_yaml(cfg) # 根据model参数决定从检测模型转换或从配置文件创建
def _from_detection_model(self, model, nc=1000, cutoff=10):
"""Creates a classification model from a YOLOv5 detection model, slicing at `cutoff` and adding a classification
layer.
"""
if isinstance(model, DetectMultiBackend): # 如果model是DetectMultiBackend实例，获取其内部模型
model = model.model # unwrap DetectMultiBackend
model.model = model.model[:cutoff] # 截断模型，保留至cutoff层作为主干网络
m = model.model[-1] # 获取截断后模型的最后一层
ch = m.conv.in_channels if hasattr(m, "conv") else m.cv1.conv.in_channels # 获取最后一层的输入通道数
c = Classify(ch, nc) # 创建分类层，传入输入通道数和类别数
c.i, c.f, c.type = m.i, m.f, "models.common.Classify" # 设置分类层的索引、来源和类型
model.model[-1] = c # 替换模型的最后一层为分类层
self.model = model.model # 将处理后的模型赋值给self.model
self.stride = model.stride # 设置步长属性
self.save = [] # 初始化保存列表
self.nc = nc # 设置类别数量属性
def _from_yaml(self, cfg):
"""Creates a YOLOv5 classification model from a specified *.yaml configuration file."""
self.model = None # 当前实现仅设置了self.model为None，实际模型构建逻辑应在后续代码中
八、定义模型结构
def parse_model(d, ch): # 定义解析YOLOv5模型结构的函数
"""从字典`d`中解析YOLOv5模型，根据输入通道数`ch`和模型架构配置各层。"""
LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}") # 打印模型概览的表头
anchors, nc, gd, gw, act, ch_mul = ( # 从配置字典中提取模型参数
d["anchors"],
d["nc"],
d["depth_multiple"],
d["width_multiple"],
d.get("activation"),
d.get("channel_multiple"),
)
if act: # 如果配置中有激活函数设置
Conv.default_act = eval(act) # 重新定义卷积层的默认激活函数
LOGGER.info(f"{colorstr('activation:')} {act}") # 打印所用的激活函数
if not ch_mul: # 如果通道乘数未设置
ch_mul = 8 # 设定默认的通道乘数值
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # 计算锚点的数量
no = na * (nc + 5) # 计算每个锚点的输出数量
layers, save, c2 = [], [], ch[-1] # 初始化模型层列表，保存列表，和最后一个通道数变量
for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]): # 遍历模型的主干和头部配置
m = eval(m) if isinstance(m, str) else m # 如果模块名是字符串，转换为对应的类对象
for j, a in enumerate(args): # 遍历模块参数
with contextlib.suppress(NameError): # 忽略NameError异常
args[j] = eval(a) if isinstance(a, str) else a # 如果参数是字符串，转换为对应的对象
n = n_ = max(round(n * gd), 1) if n > 1 else n # 计算模块的重复次数
if m in { # 判断模块类型，调整输入输出通道数和参数
Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv,
MixConv2d, Focus, CrossConv, BottleneckCSP, C3, C3TR, C3SPP,
C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x,
}:
c1, c2 = ch[f], args[0] # 获取输入和输出通道数
if c2 != no: # 如果不是输出层
c2 = make_divisible(c2 * gw, ch_mul) # 调整输出通道数
args = [c1, c2, *args[1:]] # 更新参数列表
if m in {BottleneckCSP, C3, C3TR, C3Ghost, C3x}: # 对于特定模块，插入重复次数
args.insert(2, n)
n = 1
elif m is nn.BatchNorm2d: # 如果是BatchNorm层
args = [ch[f]] # 输入通道数作为参数
elif m is Concat: # 如果是Concat层
c2 = sum(ch[x] for x in f) # 计算拼接后的通道数
elif m in {Detect, Segment}: # 如果是检测或分割层
args.append([ch[x] for x in f]) # 添加输入通道数列表
if isinstance(args[1], int): # 如果锚点数量是整数
args[1] = [list(range(args[1] * 2))] * len(f) # 转换为锚点列表
if m is Segment: # 如果是分割层
args[3] = make_divisible(args[3] * gw, ch_mul) # 调整参数
elif m is Contract: # 如果是收缩层
c2 = ch[f] * args[0] ** 2 # 计算收缩后的通道数
elif m is Expand: # 如果是扩张层
c2 = ch[f] // args[0] ** 2 # 计算扩张后的通道数
else: # 其他类型的模块
c2 = ch[f] # 输出通道数等于输入通道数
m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # 创建模块实例
t = str(m)[8:-2].replace("__main__.", "") # 获取模块类型名称
np = sum(x.numel() for x in m_.parameters()) # 计算参数数量
m_.i, m_.f, m_.type, m_.np = i, f, t, np # 附加模块的元数据
LOGGER.info(f"{i:>3}{str(f):>18}{n_:>3}{np:10.0f} {t:<40}{str(args):<30}") # 打印模块信息
save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # 更新保存列表
layers.append(m_) # 添加模块到模型层列表
if i == 0: # 如果是第一个模块
ch = [] # 清空通道数列表
ch.append(c2) # 更新通道数列表
return nn.Sequential(*layers), sorted(save) # 返回模型和排序后的保存列表
九、主函数
if __name__ == "__main__": # 当脚本直接运行时执行以下代码
parser = argparse.ArgumentParser() # 创建命令行参数解析器
parser.add_argument("--cfg", type=str, default="yolov5s.yaml", help="model.yaml") # 添加模型配置文件参数
parser.add_argument("--batch-size", type=int, default=1, help="total batch size for all GPUs") # 添加批量大小参数
parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") # 添加设备参数
parser.add_argument("--profile", action="store_true", help="profile model speed") # 添加模型速度剖析选项
parser.add_argument("--line-profile", action="store_true", help="profile model speed layer by layer") # 添加分层速度剖析选项
parser.add_argument("--test", action="store_true", help="test all yolo*.yaml") # 添加测试所有Yolo配置文件的选项
opt = parser.parse_args() # 解析命令行参数
opt.cfg = check_yaml(opt.cfg) # 检查并标准化配置文件路径
print_args(vars(opt)) # 打印解析后的参数
device = select_device(opt.device) # 选择合适的设备进行计算
# 创建模型
im = torch.rand(opt.batch_size, 3, 640, 640).to(device) # 创建随机输入张量
model = Model(opt.cfg).to(device) # 根据配置文件创建模型并移至指定设备
# 选项处理
if opt.line_profile: # 如果选择了分层剖析
model(im, profile=True) # 剖析模型的每一层
elif opt.profile: # 如果选择了整体模型剖析
results = profile(input=im, ops=[model], n=3) # 多次运行模型并收集性能数据
elif opt.test: # 如果选择了测试所有模型
for cfg in Path(ROOT / "models").rglob("yolo*.yaml"): # 遍历所有符合条件的Yolo配置文件
try:
_ = Model(cfg) # 尝试创建模型
except Exception as e: # 捕获任何异常
print(f"Error in {cfg}: {e}") # 打印错误信息
else: # 如果没有特殊选项，报告融合后的模型摘要
model.fuse() # 融合模型中的重复操作
以下代码位于yolov5/models/yolov5n.yaml
十、YOLO配置文件
# Ultralytics YOLOv5 🚀, AGPL-3.0 license
# 参数定义
nc: 80 # 类别数，即模型将识别80个不同的目标类别
depth_multiple: 0.33 # 模型深度的倍数，用于控制模型复杂度
width_multiple: 0.25 # 层通道数的倍数，用于控制模型宽度
anchors: # 锚框尺寸，用于不同尺度的特征图
- [10, 13, 16, 30, 33, 23] # 对应于P3/8特征图的锚框尺寸
- [30, 61, 62, 45, 59, 119] # 对应于P4/16特征图的锚框尺寸
- [116, 90, 156, 198, 373, 326] # 对应于P5/32特征图的锚框尺寸
# YOLOv5 v6.0 主干网络
backbone: # 主干网络定义，包含一系列的层及其参数
# [from, number, module, args] 表示从哪个层开始，重复次数，模块类型，以及参数
[
[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 卷积层，输入通道数为64，核大小为6x6，步长为2，填充为2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4 卷积层，输入通道数为128，核大小为3x3，步长为2
[-1, 3, C3, [128]], # 2-C3模块，输入通道数为128，重复3次
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8 卷积层，输入通道数为256，核大小为3x3，步长为2
[-1, 6, C3, [256]], # 4-C3模块，输入通道数为256，重复6次
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16 卷积层，输入通道数为512，核大小为3x3，步长为2
[-1, 9, C3, [512]], # 6-C3模块，输入通道数为512，重复9次
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 卷积层，输入通道数为1024，核大小为3x3，步长为2
[-1, 3, C3, [1024]], # 8-C3模块，输入通道数为1024，重复3次
[-1, 1, SPPF, [1024, 5]], # 9-SPPF模块，输入通道数为1024，核大小为5x5
]
# YOLOv5 v6.0 头部网络
head: # 头部网络定义，用于特征融合和最终预测
[
[-1, 1, Conv, [512, 1, 1]], # 卷积层，输入通道数为512，核大小为1x1
[-1, 1, nn.Upsample, [None, 2, "nearest"]], # 上采样层，放大2倍，采用最近邻插值
[[-1, 6], 1, Concat, [1]], # 拼接操作，将上采样后的特征与P4特征图拼接
[-1, 3, C3, [512, False]], # C3模块，输入通道数为512，不使用shortcut连接
[-1, 1, Conv, [256, 1, 1]], # 卷积层，输入通道数为256，核大小为1x1
[-1, 1, nn.Upsample, [None, 2, "nearest"]], # 上采样层，放大2倍，采用最近邻插值
[[-1, 4], 1, Concat, [1]], # 拼接操作，将上采样后的特征与P3特征图拼接
[-1, 3, C3, [256, False]], # C3模块，输入通道数为256，不使用shortcut连接
[-1, 1, Conv, [256, 3, 2]], # 卷积层，输入通道数为256，核大小为3x3，步长为2
[[-1, 14], 1, Concat, [1]], # 拼接操作，将特征与之前P4特征图拼接
[-1, 3, C3, [512, False]], # C3模块，输入通道数为512，不使用shortcut连接
[-1, 1, Conv, [512, 3, 2]], # 卷积层，输入通道数为512，核大小为3x3，步长为2
[[-1, 10], 1, Concat, [1]], # 拼接操作，将特征与之前P5特征图拼接
[-1, 3, C3, [1024, False]], # C3模块，输入通道数为1024，不使用shortcut连接
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect层，输入为三个不同尺度的特征图，进行目标检测
]
例：Conv卷积层位于yolov5/models/common.py
class Conv(nn.Module):
# Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
# 这是一个标准的卷积层类，接收参数包括输入通道数(ch_in)，输出通道数(ch_out)，卷积核大小(kernel)，步长(stride)，填充(padding)，组数(groups)，膨胀率(dilation)，和激活函数(activation)。
default_act = nn.SiLU() # default activation
# 设置默认的激活函数为SiLU（Swish的简化版本），这是一个自定义的类属性，可以在实例化时不改变的情况下被所有Conv类实例共享。
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
"""Initializes a standard convolution layer with optional batch normalization and activation."""
# 构造函数初始化一个标准的卷积层，带有可选的批量归一化和激活函数。
super().__init__()
# 调用父类nn.Module的构造函数初始化模块。
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
# 创建一个2D卷积层，参数为：
# c1: 输入通道数
# c2: 输出通道数
# k: 卷积核大小
# s: 步长
# autopad(k, p, d): 自动计算的padding值，如果p为None，则自动计算以保持输入输出的尺寸相同，考虑dilation的影响。
# groups: 分组卷积的组数，默认为1，表示标准卷积。
# dilation: 膨胀率，控制卷积核元素之间的间距。
# bias: 是否使用偏置项，这里设为False。
self.bn = nn.BatchNorm2d(c2)
# 创建一个2D批量归一化层，参数为c2，即卷积层的输出通道数。
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
# 设置激活函数，如果act为True，则使用默认激活函数(default_act)；
# 如果act是nn.Module的实例，则直接使用act；
# 否则，使用恒等函数Identity()。
def forward(self, x):
"""Applies a convolution followed by batch normalization and an activation function to the input tensor `x`."""
# 定义前向传播方法，对输入张量x执行卷积、批量归一化和激活函数。
return self.act(self.bn(self.conv(x)))
# 应用顺序为：卷积 -> 批量归一化 -> 激活函数。
def forward_fuse(self, x):
"""Applies a fused convolution and activation function to the input tensor `x`."""
# 定义融合的前向传播方法，只适用于不使用批量归一化的情况，直接将卷积和激活函数融合在一起。
return self.act(self.conv(x))
# 应用顺序为：卷积 -> 激活函数，省略了批量归一化步骤。