YOLOv5目标构建与损失计算

YOLOv5目标构建与损失计算

YOLOv5作为单阶段目标检测的经典算法，其高效的检测性能离不开精心设计的训练目标构建和损失计算策略。本文将深入解析YOLOv5源码中build_targets目标构建函数和ComputeLoss损失计算类的实现原理，揭开模型优化背后的关键技术。详细代码参考 YOLOv5 Github项目代码。

构建目标

以下是添加构建训练目标代码：

def build_targets(self, p, targets):
    """构建模型训练目标，从输入目标(image,class,x,y,w,h)准备损失计算所需的类别、边界框、索引和锚点

    Args:
        p (list): 模型预测输出，每个元素对应一个检测层的输出特征图
        targets (Tensor): 输入目标，形状为(nt, 6)，每行格式为(image_idx, class, x, y, w, h)

    Returns:
        tcls (list): 每个检测层对应的目标类别
        tbox (list): 每个检测层对应的目标边界框(相对于网格的xywh)
        indices (list): 每个检测层对应的(image_idx, anchor_idx, grid_y, grid_x)
        anch (list): 每个检测层对应的锚框尺寸
    """
    # 获取锚点数量和目标数量
    na, nt = self.na, targets.shape[0]  # number of anchors, targets
    tcls, tbox, indices, anch = [], [], [], []  # 初始化类别、边界框、索引和锚点列表
    
    # 归一化增益，将目标坐标从归一化形式转换到网格空间
    gain = torch.ones(7, device=self.device)  # 7维对应(image_idx, class, x, y, w, h, anchor_idx)
    
    # 创建锚点索引，形状(na, nt)，用于标识每个目标对应的锚点
    ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt)
    
    # 将目标重复na次，并添加锚点索引，形状变为(na, nt, 7)
    targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None]), 2)

    # 设置网格偏移参数
    g = 0.5  # 偏移量阈值，用于中心点偏移判断
    off = torch.tensor(  # 定义5种偏移量（中心+四个方向）
        [
            [0, 0],    # 中心
            [1, 0],     # 右
            [0, 1],     # 下
            [-1, 0],    # 左
            [0, -1],    # 上
        ],
        device=self.device,
    ).float() * g  # 应用偏移系数

    # 遍历每个检测层（不同尺度的特征图）
    for i in range(self.nl):
        # 获取当前层的锚点尺寸和特征图形状
        anchors = self.anchors[i]  # 当前层锚点尺寸，形状(na, 2)
        shape = p[i].shape  # 预测特征图形状(batch_size, anchors, grid_y, grid_x, params)
        
        # 设置归一化增益（将xywh转换到当前特征图尺度）
        gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]]  # xyxy增益为特征图的宽高

        # 将目标坐标转换到当前特征图尺度
        t = targets * gain  # 形状(na, nt, 7)
        
        if nt:  # 存在目标时处理
            # 计算目标宽高与锚点宽高的比例
            r = t[..., 4:6] / anchors[:, None]  # wh比例，形状(na, nt, 2)
            
            # 筛选满足宽高比例阈值的锚点（最大比例小于hyp['anchor_t']）
            j = torch.max(r, 1 / r).max(2)[0] < self.hyp["anchor_t"]  # 形状(na, nt)
            t = t[j]  # 过滤后的目标，形状(nt1, 7)

            # 计算网格偏移量
            gxy = t[:, 2:4]  # 目标在特征图上的xy坐标，形状(nt1, 2)
            gxi = gain[[2, 3]] - gxy  # 反向坐标（用于边界判断）
            
            # 生成偏移掩码（判断是否需要向相邻网格分配目标）
            j, k = ((gxy % 1 < g) & (gxy > 1)).T  # 右、下方向偏移条件
            l, m = ((gxi % 1 < g) & (gxi > 1)).T  # 左、上方向偏移条件
            j = torch.stack((torch.ones_like(j), j, k, l, m))  # 合并所有条件，形状(5, nt1)
            
            # 扩展目标到5个偏移位置（中心+四个方向）
            t = t.repeat((5, 1, 1))[j]  # 形状(5, nt1, 7) -> (nt2, 7)
            offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]  # 对应偏移量，形状(nt2, 2)
        else:  # 无目标时处理
            t = targets[0]  # 取空目标
            offsets = 0  # 无偏移

        # 解包处理后的目标数据
        bc = t[:, :2]  # (image_idx, class)
        gxy = t[:, 2:4]  # 网格xy坐标
        gwh = t[:, 4:6]  # 网格wh尺寸
        a = t[:, 6]  # 锚点索引
        
        # 转换数据类型并拆分
        a = a.long().view(-1)  # 锚点索引转为长整型
        (b, c) = bc.long().T  # 图像索引和类别

        # 计算目标所在的网格坐标（考虑偏移）
        gij = (gxy - offsets).long()
        gi, gj = gij.T  # 分解为x,y坐标

        # 将网格坐标限制在特征图范围内
        gj = gj.clamp_(0, shape[2] - 1)
        gi = gi.clamp_(0, shape[3] - 1)

        # 存储当前层的信息
        indices.append((b, a, gj, gi))  # 图像索引、锚点索引、网格y,x
        tbox.append(torch.cat((gxy - gij, gwh), 1))  # 相对于网格的xy和原始wh
        anch.append(anchors[a])  # 对应的锚点尺寸
        tcls.append(c)  # 目标类别

    return tcls, tbox, indices, anch

关键步骤解析：

1. 锚点匹配
   通过计算目标宽高与锚点宽高的比例，筛选出宽高比小于阈值anchor_t的锚点。这确保目标被分配到最合适尺寸的锚点。

2. 网格偏移处理
   当目标中心靠近网格边界时（偏移量g=0.5），将目标分配给相邻的网格。这增加了正样本数量，有助于模型学习。

3. 多尺度分配
   不同检测层（不同特征图尺度）处理不同大小的目标。通过gain将归一化坐标转换到对应特征图尺度，实现多尺度训练。

4. 数据格式转换
   最终输出的tbox存储相对于网格单元的坐标偏移和原始宽高，用于计算定位损失。indices则记录目标对应的位置信息，用于从预测结果中提取对应预测值。

该函数核心思想是将每个目标分配到最合适的特征图层、网格位置和锚点尺寸，同时考虑中心点偏移以增加匹配机会，最终构建用于计算分类和定位损失的训练目标。

计算损失

以下是添加详细注释后的YOLOv5损失计算代码：

class ComputeLoss:
    """计算YOLOv5模型的总损失，包含分类损失、边界框损失和置信度损失"""

    sort_obj_iou = False  # 是否对目标IoU进行排序（默认关闭）

    def __init__(self, model, autobalance=False):
        """初始化损失计算模块
        Args:
            model: 要计算损失的模型
            autobalance: 是否自动平衡各检测层的损失权重
        """
        device = next(model.parameters()).device  # 获取模型所在设备
        h = model.hyp  # 获取超参数配置

        # 定义基础损失函数
        BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["cls_pw"]], device=device))  # 分类损失
        BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["obj_pw"]], device=device))  # 置信度损失

        # 标签平滑参数（正样本和负样本的平滑系数）
        self.cp, self.cn = smooth_BCE(eps=h.get("label_smoothing", 0.0))

        # Focal Loss配置（如果gamma>0则启用）
        g = h["fl_gamma"]
        if g > 0:
            BCEcls = FocalLoss(BCEcls, g)  # 分类Focal Loss
            BCEobj = FocalLoss(BCEobj, g)  # 置信度Focal Loss

        # 获取模型Detect层
        m = de_parallel(model).model[-1]  # 获取最后一个模块（Detect）
        
        # 设置各检测层的损失平衡系数（不同尺度的特征图赋予不同权重）
        self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02])  # P3-P7的默认系数
        self.ssi = list(m.stride).index(16) if autobalance else 0  # 用于自动平衡的参考层（stride=16的层）
        
        # 存储关键参数
        self.BCEcls = BCEcls
        self.BCEobj = BCEobj
        self.gr = 1.0  # IoU比例系数（用于混合标签）
        self.hyp = h
        self.autobalance = autobalance
        self.na = m.na  # 每层的锚点数量
        self.nc = m.nc  # 类别数量
        self.nl = m.nl  # 检测层数量
        self.anchors = m.anchors  # 锚点尺寸
        self.device = device

    def __call__(self, p, targets):
        """计算总损失
        Args:
            p: 模型预测输出列表，每个元素对应一个检测层的预测结果
            targets: 真实标签张量，形状为(nt, 6)，每行格式为(image_idx, class, x, y, w, h)
        
        Returns:
            (总损失, 各损失分量) 元组
        """
        # 初始化各损失分量
        lcls = torch.zeros(1, device=self.device)  # 分类损失
        lbox = torch.zeros(1, device=self.device)  # 边界框损失
        lobj = torch.zeros(1, device=self.device)  # 置信度损失
        
        # 构建训练目标（关键步骤）
        tcls, tbox, indices, anchors = self.build_targets(p, targets)  # 获取匹配后的目标

        # 遍历每个检测层计算损失
        for i, pi in enumerate(p):  # i: 层索引, pi: 该层预测结果
            b, a, gj, gi = indices[i]  # 分解匹配结果：
                                       # b: 图片索引, a: 锚点索引
                                       # gj, gi: 网格y,x坐标
                                       
            tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device)  # 初始化目标置信度张量

            n = b.shape[0]  # 当前层的目标数量
            if n:
                # 分解预测结果（使用split替代新版tensor_split以兼容旧版本）
                pxy, pwh, _, pcls = pi[b, a, gj, gi].split((2, 2, 1, self.nc), 1)
                
                # --------------------- 边界框回归损失计算 ---------------------
                # 解码预测框坐标（基于YOLOv5的改进解码方式）
                pxy = pxy.sigmoid() * 2 - 0.5  # 将xy预测值从(0,1)映射到(-0.5,1.5)
                pwh = (pwh.sigmoid() * 2) ** 2 * anchors[i]  # 将wh预测值从(0,4)映射到(0,4*anchor)
                pbox = torch.cat((pxy, pwh), 1)  # 组合成完整预测框(xywh格式)
                
                # 计算CIoU损失
                iou = bbox_iou(pbox, tbox[i], CIoU=True).squeeze()  # 形状(n,)
                lbox += (1.0 - iou).mean()  # 平均IoU损失

                # --------------------- 置信度目标生成 ---------------------
                iou = iou.detach().clamp(0).type(tobj.dtype)  # 分离计算图并确保非负
                if self.sort_obj_iou:  # 按IoU排序（可选）
                    j = iou.argsort()
                    b, a, gj, gi, iou = b[j], a[j], gj[j], gi[j], iou[j]
                if self.gr < 1:  # 混合真实标签和预测置信度（当gr=1时完全使用预测值）
                    iou = (1.0 - self.gr) + self.gr * iou
                tobj[b, a, gj, gi] = iou  # 将IoU作为置信度目标

                # --------------------- 分类损失计算 ---------------------
                if self.nc > 1:  # 仅当类别数>1时计算分类损失
                    t = torch.full_like(pcls, self.cn, device=self.device)  # 初始化目标为负样本
                    t[range(n), tcls[i]] = self.cp  # 设置正样本位置
                    lcls += self.BCEcls(pcls, t)  # 计算分类BCE损失

            # --------------------- 置信度损失计算 ---------------------
            obji = self.BCEobj(pi[..., 4], tobj)  # 置信度损失（pi[...,4]是原始预测值）
            lobj += obji * self.balance[i]  # 加权后的置信度损失
            if self.autobalance:  # 自动平衡各层损失权重
                self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()

        # --------------------- 损失加权与整合 ---------------------
        if self.autobalance:  # 归一化平衡系数
            self.balance = [x / self.balance[self.ssi] for x in self.balance]
        lbox *= self.hyp["box"]  # 边界框损失加权
        lobj *= self.hyp["obj"]  # 置信度损失加权
        lcls *= self.hyp["cls"]  # 分类损失加权
        bs = tobj.shape[0]  # 获取batch size

        # 返回总损失和各损失分量（总损失乘以batch size）
        return (lbox + lobj + lcls) * bs, torch.cat((lbox, lobj, lcls)).detach()

关键实现细节解析

1. 预测框解码

   • XY坐标：通过sigmoid缩放至(0,1)后，乘以2减0.5，将中心点范围从网格中心的±0.5扩展到相邻网格（-0.5到1.5），增强对小目标的检测能力
   • WH尺寸：使用sigmoid的(0,4)次方缩放，保证预测框尺寸不超过4倍锚框尺寸，避免梯度爆炸

2. 损失自动平衡

   • 通过balance数组为不同检测层分配不同权重（浅层特征权重更高）
   • 当autobalance=True时，根据各层损失动态调整权重，使各层损失贡献均衡

3. 置信度目标生成

   • 使用预测框与真实框的IoU作为监督信号（tobj），而非固定1.0
   • 引入gr参数（梯度比率）实现标签平滑：iou = (1.0 - gr) + gr * iou

4. 分类标签平滑

   • 正样本标签值设为cp（如0.95），负样本设为cn（如0.05）
   • 缓解类别不平衡问题，防止模型过度自信

5. 多尺度训练策略

   • 不同检测层（P3-P5或P3-P7）处理不同尺度的目标
   • 通过balance参数平衡浅层（小目标）和深层（大目标）的损失贡献

各损失分量说明

该实现通过动态目标分配、多尺度损失平衡和先进的IoU计算方式，有效提升了YOLOv5的检测性能。