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

YOLOv5核心代码深度解析

news 2025/10/28 9:43:16

引言

目标检测是计算机视觉领域的重要任务，而 YOLO（You Only Look Once）系列模型凭借其高效的实时检测能力备受关注。本文将深入解析 YOLOv5 的核心代码，包括模型结构定义、前向传播过程和训练流程，帮助读者理解其工作原理。

模型核心组件解析

1. 检测层（Detect 类）

Detect 类是 YOLOv5 的输出层，负责生成最终的检测结果。其核心代码：

class Detect(nn.Module):stride = None  #  strides computed during buildexport = False  #  onnx exportdef __init__(self, nc=80, anchors=(), ch=()):  # detection layersuper(Detect, self).__init__()self.nc = nc  # 类别数量self.no = nc + 5  # 每个锚点的输出数量 (x,y,w,h,conf+classes)self.nl = len(anchors)  # 检测层数self.na = len(anchors[0]) // 2  # 每个检测层的锚点数量self.grid = [torch.zeros(1)] * self.nl  # 初始化网格a = torch.tensor(anchors).float().view(self.nl, -1, 2)self.register_buffer('anchors', a)  # 注册锚点为缓冲区 (nl,na,2)self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2))  # 锚点网格self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # 输出卷积层

在初始化时，Detect 类主要完成：

定义类别数量和输出维度
处理锚点信息并注册为缓冲区（不参与梯度计算）
创建卷积层用于生成最终检测结果

2. 前向传播过程

Detect 类的 forward 方法实现了从特征图到检测框的转换：

def forward(self, x):z = []  # 推理输出self.training |= self.exportfor i in range(self.nl):x[i] = self.m[i](x[i])  # 卷积操作bs, _, ny, nx = x[i].shape  # 形状转换 (bs,255,20,20) -> (bs,3,20,20,85)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.grid[i].shape[2:4] != x[i].shape[2:4]:self.grid[i] = self._make_grid(nx, ny).to(x[i].device)y = x[i].sigmoid()# 计算边界框坐标y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i]  # xyy[..., 2:4] = (y[..., 2:4] * 2) **2 * self.anchor_grid[i]  # whz.append(y.view(bs, -1, self.no))return x if self.training else (torch.cat(z, 1), x)

前向传播过程的核心是：

对每个尺度的特征图进行卷积操作
转换特征图形状以匹配锚点结构
在推理阶段，将模型输出转换为实际图像坐标的边界框

3. 基础网络组件

YOLOv5 使用了多种基础组件构建网络，定义：

Bottleneck 模块：

class BottleneckCSP(nn.Module):# CSP瓶颈结构 https://github.com/WongKinYiu/CrossStagePartialNetworksdef __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):super(BottleneckCSP, self).__init__()c_ = int(c2 * e)  # 隐藏层通道数self.cv1 = Conv(c1, c_, 1, 1)self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)self.cv4 = Conv(2 * c_, c2, 1, 1)self.bn = nn.BatchNorm2d(2 * c_)  # 应用于cat(cv2, cv3)self.act = nn.LeakyReLU(0.1, inplace=True)self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])def forward(self, x):y1 = self.cv3(self.m(self.cv1(x)))y2 = self.cv2(x)return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))

Bottleneck 模块采用 1x1 卷积降维再 3x3 卷积升维的结构，当 shortcut=True 且输入输出通道相同时会添加残差连接。

BottleneckCSP 模块：

class BottleneckCSP(nn.Module):# CSP瓶颈结构 https://github.com/WongKinYiu/CrossStagePartialNetworksdef __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):super(BottleneckCSP, self).__init__()c_ = int(c2 * e)  # 隐藏层通道数self.cv1 = Conv(c1, c_, 1, 1)self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)self.cv4 = Conv(2 * c_, c2, 1, 1)self.bn = nn.BatchNorm2d(2 * c_)  # 应用于cat(cv2, cv3)self.act = nn.LeakyReLU(0.1, inplace=True)self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])def forward(self, x):y1 = self.cv3(self.m(self.cv1(x)))y2 = self.cv2(x)return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))

CSP 结构将输入分为两部分，一部分经过多个 Bottleneck，另一部分直接连接，最后将两部分拼接，有效减少了计算量同时保持精度。

Focus 模块：

class Focus(nn.Module):# 将宽高信息集中到通道维度def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):super(Focus, self).__init__()self.conv = Conv(c1 * 4, c2, k, s, p, g, act)def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))

Focus 模块通过对输入图像进行隔点采样并拼接，在减少空间维度的同时增加通道数，有效保留了原始图像信息。

模型构建过程

YOLOv5 通过配置文件定义网络结构，然后通过parse_model函数解析配置并构建模型：

def parse_model(d, ch):  # model_dict, input_channels(3)logger.info('\n%3s%18s%3s%10s  %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments'))anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # 锚点数量no = na * (nc + 5)  # 输出数量 = 锚点 * (类别 + 5)layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch outfor i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  # from, number, module, argsm = eval(m) if isinstance(m, str) else m  # 解析模块for j, a in enumerate(args):try:args[j] = eval(a) if isinstance(a, str) else a  # 解析参数except:passn = max(round(n * gd), 1) if n > 1 else n  # 深度增益if m in [Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]:c1, c2 = ch[f], args[0]c2 = make_divisible(c2 * gw, 8) if c2 != no else c2args = [c1, c2, *args[1:]]if m in [BottleneckCSP, C3]:args.insert(2, n)n = 1# ... 其他模块处理m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args)  # 构建模块# ... 记录和保存信息layers.append(m_)ch.append(c2)return nn.Sequential(*layers), sorted(save)

parse_model函数的主要工作：

解析配置文件中的网络结构定义
根据深度倍数 (gd) 和宽度倍数 (gw) 调整网络深度和宽度
构建各个网络层并连接成完整模型

训练流程解析

训练过程定义在train函数中，主要包括以下步骤：

1. 初始化与配置

def train(hyp, opt, device, tb_writer=None):logger.info(f'Hyperparameters {hyp}')log_dir = Path(tb_writer.log_dir) if tb_writer else Path(opt.logdir) / 'evolve'  # 日志目录wdir = log_dir / 'weights'  # 权重目录os.makedirs(wdir, exist_ok=True)last = wdir / 'last.pt'best = wdir / 'best.pt'results_file = str(log_dir / 'results.txt')# 加载数据配置with open(opt.data) as f:data_dict = yaml.load(f, Loader=yaml.FullLoader)  # 数据字典train_path = data_dict['train']test_path = data_dict['val']nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])

训练开始前会进行必要的初始化：

设置日志和权重保存目录
加载数据配置文件
解析类别信息

2. 模型加载与初始化

# 模型加载
pretrained = weights.endswith('.pt')
if pretrained:ckpt = torch.load(weights, map_location=device)  # 加载检查点model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc).to(device)  # 创建模型exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else []  # 排除的键state_dict = ckpt['model'].float().state_dict()  # 转换为FP32state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude)  # 交集model.load_state_dict(state_dict, strict=False)  # 加载权重
else:model = Model(opt.cfg, ch=3, nc=nc).to(device)  # 创建新模型

支持加载预训练模型进行迁移学习，通过intersect_dicts函数确保权重与模型结构匹配。

3. 优化器与学习率调度

# 优化器参数组
pg0, pg1, pg2 = [], [], []  # 优化器参数组
for k, v in model.named_parameters():v.requires_grad = Trueif '.bias' in k:pg2.append(v)  # 偏置elif '.weight' in k and '.bn' not in k:pg1.append(v)  # 应用权重衰减else:pg0.append(v)  # 其他参数# 优化器选择
if opt.adam:optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))
else:optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)# 学习率调度器
lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp['lrf']) + hyp['lrf']  # 余弦退火
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)

YOLOv5 将参数分为三组进行优化，并使用余弦退火调度学习率，有助于模型更好地收敛。

4. 训练循环

for epoch in range(start_epoch, epochs):  #  epoch循环model.train()mloss = torch.zeros(4, device=device)  # 平均损失pbar = enumerate(dataloader)logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'targets', 'img_size'))if rank in [-1, 0]:pbar = tqdm(pbar, total=nb)  # 进度条optimizer.zero_grad()for i, (imgs, targets, paths, _) in pbar:  #  batch循环ni = i + nb * epoch  # 累计批次imgs = imgs.to(device, non_blocking=True).float() / 255.0  # 归一化# 热身训练if ni <= nw:# 调整学习率和动量xi = [0, nw]accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round())# ...# 多尺度训练if opt.multi_scale:sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs  # 随机尺寸# ...# 前向传播with amp.autocast(enabled=cuda):pred = model(imgs)  # 前向传播loss, loss_items = compute_loss(pred, targets.to(device), model)  # 计算损失# 反向传播scaler.scale(loss).backward()# 优化步骤if ni % accumulate == 0:scaler.step(optimizer)  # 优化器步骤scaler.update()optimizer.zero_grad()if ema:ema.update(model)# 日志输出# ...# 学习率调度scheduler.step()# 评估与保存# ...

训练循环的核心机制包括：