代码层面上学习yolo12

Chirp
193 阅读14分钟

总览

对 YOLO 的具体运作原理仍然不甚理解，来读下 ultralytics 库的代码吧。调试 YOLO v12 的推理，以及训练。

写完后就感觉，还不如直接用 ultralytics 库来推理或者训练 YOLO 模型。现在的 YOLO 不论是网络结构还是训练方法都整得太复杂了。

本文写得太乱了。现在的 YOLO 用上了数不清的技巧和 trick，头快炸了。

感觉最大的收获是了解到了 CIoU（Complete IoU）和 DFL（Distribution Focal）的概念。

推理流程

  • 图像预处理
    • 使用 LetterBox，缩放与填充为 最长边不超过 640、边长为 32 倍数的图片
    • 除以 255，数值归一化到 [0, 1]
  • 推理
    • 进入到 PyTorch 后端 BaseModel，继承于 nn.Module
    • 获得 [1, 84, 3360] 和 [1, 144, 32, 80]、[1, 144, 16, 40]、[1, 144, 8, 20]
  • 后处理

模型推理流程

model = [
    Conv,  # 3 -> 64, kernel_size=3, stride=2, padding=1
    Conv,  # 64 -> 128, kernel_size=3, stride=2, padding=1
    C3k2
    Conv,  # 256 -> 256, kernel_size=3, stride=2, padding=1
    C3k2
    Conv,  # 512 -> 512, kernel_size=3, stride=2, padding=1
    A2C2f,  # attention
    Conv,  # 512 -> 512, kernel_size=3, stride=2, padding=1
    A2C2f,  # attention
    Upsample(scale_factor=2.0, mode='nearest'),
    Concat(-1, 6),  # 拼接上层和第 6 层的 x 结果
    A2C2f,
    Upsample(scale_factor=2.0, mode='nearest'),
    Concat(-1, 4),  # 拼接上层和第 4 层的 x 结果
    A2C2f,
    Conv,  # 256 -> 256, kernel_size=3, stride=2, padding=1
    Concat(-1, 11),  # 拼接上层和第 11 层的 x 结果
    A2C2f,
    Conv,  # 512 -> 512, kernel_size=3, stride=2, padding=1
    Concat(-1, 8),  # 拼接上层和第 8 层的 x 结果
    C3k2
    Detect  # 接受第 14, 17, 20 层的 x 结果，记为 x_1, x_2, x_3
]

模型详情

{
    10: [-1, 6],
    13: [-1, 4],
    16: [-1, 11],
    19: [-1, 8],
    21: [14, 17, 20],
}

尝试伪代码，以类名代表网络层。请结合注释阅读。

没有标注 padding 的其值为 0。

# x 是输入张量

# 定义会频繁用到的层
Conv = [
    Conv2d,
    BatchNorm2d,
    SiLU,
]
Bottleneck = [
    Conv,
    Conv,
]
# C3k 继承于 C3
C3k = [
    cat(  # 在通道维度拼接
        [
            Conv,  # 2n -> n, kernel_size=1, stride=1
            Bottleneck,  # n -> n, kernel_size=3, stride=1, padding=1
            Bottleneck,  # n -> n, kernel_size=3, stride=1, padding=1
        ],
        Conv,  # 2n -> n, kernel_size=1, stride=1
    ),
    Conv,  # 2n -> 2n, kernel_size=1, stride=1
]

# 模型推理流程
model = [
    Conv(x),  # 3 -> 64, kernel_size=3, stride=2, padding=1
    Conv(x),  # 64 -> 128, kernel_size=3, stride=2, padding=1
    C3k2[  # 继承于 C2f
        Conv(x),  # 128 -> 128, kernel_size=1, stride=1
        y = chunk(x),  # 在通道维度平分出两份 tensor
        y_1 = C3k(y[-1]),  # n=32
        y_2 = C3k(y[-1]),  # n=32
        y.extend([y_1, y_2]), 
        x = cat(y),  # 在通道维度拼接
        Conv(x),  # 256 -> 256, kernel_size=1, stride=1
    ],
    Conv(x),  # 256 -> 256, kernel_size=3, stride=2, padding=1
    C3k2[
        Conv(x),  # 256 -> 256, kernel_size=1, stride=1
        y = chunk(x),  # 在通道维度平分出两份 tensor
        y_1 = C3k(y[-1]),  # n=64
        y_2 = C3k(y[-1]),  # n=64
        y.extend([y_1, y_2]), 
        x = cat(y),  # 在通道维度拼接
        Conv(x),  # 512 -> 512, kernel_size=1, stride=1
    ],
    Conv(x),  # 512 -> 512, kernel_size=3, stride=2, padding=1
    A2C2f[
        y = x
        Conv(y),  # 512 -> 256, kernel_size=1, stride=1
        y_1 = 2 * ABlock,
        y_2 = 2 * ABlock,
        y_3 = 2 * ABlock,
        y_4 = 2 * ABlock,
        y = cat(y, y_1, y_2, y_3, y_4),  # 在通道维度拼接
        Conv(y),  # 1280 -> 512, kernel_size=1, stride=1
        x = x + y * gamma,  # 通道层面的残差控制
    ],
    Conv(x),  # 512 -> 512, kernel_size=3, stride=2, padding=1
    A2C2f[
        y = x
        Conv(y),  # 512 -> 256, kernel_size=1, stride=1
        y_1 = 2 * ABlock,
        y_2 = 2 * ABlock,
        y_3 = 2 * ABlock,
        y_4 = 2 * ABlock,
        y = cat(y, y_1, y_2, y_3, y_4),  # 在通道维度拼接
        Conv(y),  # 1280 -> 512, kernel_size=1, stride=1
        x = x + y * gamma,  # 通道层面的残差控制
    ],
    Upsample(scale_factor=2.0, mode='nearest'),
    Concat(-1, 6),  # 拼接上层和第 6 层的 x 结果
    A2C2f[
        y = x
        Conv(y),  # 1024 -> 256, kernel_size=1, stride=1
        y_1 = C3k(y[-1]),  # n=128
        y_2 = C3k(y[-1]),  # n=128
        y = cat(y, y_1, y_2),  # 在通道维度拼接
        Conv(y),  # 768 -> 512, kernel_size=1, stride=1
    ],
    Upsample(scale_factor=2.0, mode='nearest'),
    Concat(-1, 4),  # 拼接上层和第 4 层的 x 结果
    A2C2f[
        y = x
        Conv(y),  # 1024 -> 128, kernel_size=1, stride=1
        y_1 = C3k(y[-1]),  # n=64
        y_2 = C3k(y[-1]),  # n=64
        y = cat(y, y_1, y_2),  # 在通道维度拼接
        Conv(y),  # 384 -> 256, kernel_size=1, stride=1
    ],
    Conv(x),  # 256 -> 256, kernel_size=3, stride=2, padding=1
    Concat(-1, 11),  # 拼接上层和第 11 层的 x 结果
    A2C2f[
        y = x
        Conv(y),  # 768 -> 256, kernel_size=1, stride=1
        y_1 = C3k(y[-1]),  # n=128
        y_2 = C3k(y[-1]),  # n=128
        y = cat(y, y_1, y_2),  # 在通道维度拼接
        Conv(y),  # 768 -> 512, kernel_size=1, stride=1
    ],
    Conv(x),  # 512 -> 512, kernel_size=3, stride=2, padding=1
    Concat(-1, 8),  # 拼接上层和第 8 层的 x 结果
    C3k2[
        Conv(x),  # 1024 -> 512, kernel_size=1, stride=1
        y = chunk(x),  # 在通道维度平分出两份 tensor
        y_1 = C3k(y[-1]),  # n=128
        y_2 = C3k(y[-1]),  # n=128
        y.extend([y_1, y_2]), 
        x = cat(y),  # 在通道维度拼接
        Conv(x),  # 1024 -> 512, kernel_size=1, stride=1
    ],
    Detect[  # 接受第 14, 17, 20 层的 x 结果，记为 x_1, x_2, x_3
        x_1_a = x_1,
        x_1_b = x_1,
        Conv(x_1_a),  # 256 -> 64, kernel_size=3, stride=1, padding=1
        Conv(x_1_a),  # 64 -> 64, kernel_size=3, stride=1, padding=1
        Conv2d(x_1_a),  # 64 -> 64, kernel_size=1, stride=1
        2 * [
            Conv(x_1_b),  # 256 -> 256, kernel_size=3, stride=1, padding=1, groups=256
            Conv(x_1_b),  # 256 -> 256, kernel_size=1, stride=1
        ]
        Conv2d(x_1_b),  # 256 -> 80, kernel_size=1, stride=1
        x_1 = cat(x_1_a, x_1_b),  # 在通道维度拼接

        x_2_a = x_2,
        x_2_b = x_2,
        Conv(x_2_a),  # 512 -> 64, kernel_size=3, stride=1, padding=1
        Conv(x_2_a),  # 64 -> 64, kernel_size=3, stride=1, padding=1
        Conv2d(x_2_a),  # 64 -> 64, kernel_size=1, stride=1
        2 * [
            Conv(x_2_b),  # 512 -> 512, kernel_size=3, stride=1, padding=1, groups=512
            Conv(x_2_b),  # 512 -> 256, kernel_size=1, stride=1
        ]
        Conv2d(x_2_b),  # 256 -> 80, kernel_size=1, stride=1
        x_2 = cat(x_2_a, x_2_b),  # 在通道维度拼接

        x_3_a = x_3,
        x_3_b = x_3,
        Conv(x_3_a),  # 512 -> 64, kernel_size=3, stride=1, padding=1
        Conv(x_3_a),  # 64 -> 64, kernel_size=3, stride=1, padding=1
        Conv2d(x_3_a),  # 64 -> 64, kernel_size=1, stride=1
        2 * [
            Conv(x_3_b),  # 512 -> 512, kernel_size=3, stride=1, padding=1, groups=512
            Conv(x_3_b),  # 512 -> 256, kernel_size=1, stride=1
        ]
        Conv2d(x_3_b),  # 256 -> 80, kernel_size=1, stride=1
        x_3 = cat(x_3_a, x_3_b),  # 在通道维度拼接

        x_cat = cat(x_1, x_2, x_3),  # WH 拍平后在 WH 维度进行拼接
        box, cls = x_cat.split(64, 80),  # 在通道维度拆分。对应刚才的 ab 处理

        view(box),  # 在通道维度拆分，64 -> (4, 16)
        transpose(box),  # (4, 16) -> (16, 4)。此时 16 变成了通道，4 变成了 H
        softmax(box),  # 施加在通道维度
        Conv2d(box),  # 16 -> 1 且 bias=False 的 1×1 卷积。权重手动指定为了 arange(16)

        sigmoid(cls)

        # 此时的 box 已经具有 xywh 方形框形式。每个小格子都有一个框
        # 后续处理围绕着偏移处理。就不写了
    ]
]

ABlock，会先后经过 x = x + self.attn(x) 和 x = x + self.mlp(x)，和 Transformer 很像。具体来说：

  • self.attn，AAttn
    • 维度 256，头数 8，HW 维度摊平作为序列长度
    • 标准 scaled-dot-product-attention，没有 mask，没有位置编码
    • 最后额外有 x = x + Conv(v)，kernel_size=7, stride=1, padding=3, act=False，加上 v 的值
    • 最后额外有 x = Conv(x)，kernel_size=1, stride=1, act=False，作为 proj
  • self.mlp，
    • Conv，256 -> 307, kernel_size=1, stride=1
    • Conv，307 -> 256, kernel_size=1, stride=1, act=False

可见，attention 在最后让 x 加上了 value 的值，mlp 是中间维度不是很高的两次卷积实现的。

后处理流程

现有一个维度为 [batch, 84, 3360] 的 tensor。其中 3360 意味着不同分割框（随输入图像大小变化而改变），84 中分别含有 4 份 xyhw 信息代表标注框、80 份分类信息。

变量 xc 指向 3360 个框中 最大的分类概率大于 conf_thres 值的框。这样一来，绝大部分的框都被淘汰了。

会用到 nms（non-maximum suppression） 算法进一步去重筛选。算法步骤：

  1. 输入三个参数：box 框、各个框的置信度分数、目标 iou 阈值
  2. 将置信度最高的框转移到保留队列
  3. 剩余的框计算其与上一步被转移的框的 iou，去掉大于目标 iou 阈值的框
  4. 回到 2，不断重复，直到无框可移

不同类别的框不应该相互覆盖，影响 iou 计算。按理说应该一个类别一个类别地应用 nms 来筛选框。

源代码用到了一个很讨巧的方法，使得不同类别的 box 框能在一个批次里同时被处理，缩短时间。

方法是：让每个类别的框的坐标加上不同偏移。具体来说，若是有 4 个类别，输入图像最大边长为 7680，那么让每个类别的框坐标相互间隔 7680、7680 * 2、7680 * 3 即可。

训练过程（v8DetectionLoss）

代码将 loss 的计算封装在 v8DetectionLoss。输入模型预测结果和标签，进行一些预处理转换得到：

  • pd_scores，predicted scores，预测类别置信度
    • 维度 [batch, seq, num_cls]
  • pd_bboxes，predicted bboxes，预测框
    • 维度 [batch, seq, 4]
  • anc_points，anchor points。各个分区的中心点
    • 维度 [seq, 2]
  • gt_labels，ground truth labels，标注框对应类别
    • 维度 [batch, max_seq_len, num_cls]
  • gt_bboxes，ground truth bboxes，标注框
    • 维度 [batch, max_seq_len, 4]
  • mask_gt，由于需要填充到 max_seq_len 故有此 mask
    • 维度 [batch, max_seq_len, 1]

将这些参数输入到 self.assigner（TaskAlignedAssigner），获得调整过的学习目标。

损失分为三个部分，

  • iou loss，使用的 CIoU 计算 bbox损失
  • cls loss，用 BCEWithLogitsLoss 计算置信度损失
  • dfl loss，使用 DFL 计算 bbox 损失。

由于用到了 softmax 计算置信度，使用 MSELoss 会导致在 0 和 1 处梯度很小。置信度损失使用 BCELoss 会更合适。

cls loss 和 dfl loss 都会用目标置信度作为权重进行加权相加。三种损失都会用目标置信度之和 进行归一化。

BboxLoss 计算出 dfl loss

模型没有直接输出 4 个特征来指示框的 xyhw，而是把输出特征翻了 16 倍。这是因为模型把连续的数值预测问题转换为了离散的分类问题，从 0 到 15 总共 16 个分类。换句话说，模型用 16 个数字来代表范围在 [0, 15] 的一个数字。

dfl loss 的计算逻辑封装在 BboxLoss 类内。DFL（Distribution Focal）将连续坐标回归问题转换为离散的分类问题，计算时用到了两次 F.cross_entropy 以考虑小数取整时的向上和向下取整两个情况。

模型用的分区，最小有 8×8，最大有 32×32。只有 16 个分类就意味着最大输出为 15 倍分区边长，那么模型能画出的最大框就为 480×480。

TaskAlignedAssigner 获得调整过的学习目标

self.get_pos_mask() 筛选分区，减轻计算量。

self.select_highest_overlaps()，在上一步 mask 挑选的基础上，如果一个分区同时预测了多个预测框，用 overlaps 决定到底用哪一个预测框。

self.get_targets()，对于 mask 挑选的结果，将 gt_labels 和 gt_bboxes 转换为合适的形式，获得学习目标 target_bboxes target_scores，便于 loss 计算。

Normalize，该步骤会调整上一步获得的仅含 0 值 和 1 值的 target_scores。思路是“质量”越高的分区目标分数也越高，而不是单纯的 0 和 1。这个“质量”追根溯源是通过置信度确定的。

至此，获得了三个重要变量：

  • target_bboxes，
  • target_scores
  • fg_mask

self.get_pos_mask() 具体如何筛选分区

不可能对所有分区都计算 loss，否则计算量太大了。代码在 self.get_pos_mask() 中按顺序用到以下方法筛选分区。

select_candidates_in_gts()，输入一系列中心点和标注框，返回一个用于指示分区中心点是否在框内的 mask mask_in_gts。

get_box_metrics()，计算获得预测框与标注框的重叠率 overlaps，以及重叠率与对应类别置信度的乘积指标 align_metric。

  • 会先用 mask_in_gts * mask_gt 减少计算量
  • 重叠率用 CIoU 衡量
  • 虽然单个分区只会预测出一个框，但此时假设分区会预测所有标注框
  • 用 mask 跳过中心点不在标注框内的分区
  • align_metric 计算方式：置信度^alpha * 重叠率。这个 alpha 我调试出来被设为了 0.5

select_topk_candidates()，输入刚刚算出的重叠率指标 align_metric，输出进一步筛选的分区，返回为 mask mask_topk。对于单张图：

  1. 对图中所有标注框，分别选出指标最大的 self.topk 个对象（10 个）
    1. 意味着，总共选出 (num_cls, 10) 个分区的指标
  2. 统计分区被选中的次数
  3. 若有分区被选中不止一次（说明这个分区预测了不止一个框），则视为一次都没被选中
  4. 返回被选中且只被选中一次的分区

最终，用 mask_topk * mask_in_gts * mask_gt 结果作为筛出分区的 mask。八千个分区就筛得只剩下三十几个了。

一些值得注意的地方

模型输出了 4 * 16 个特征来指示预测框

模型没有直接输出 4 个特征来指示框的 xyhw，而是把输出特征翻了 16 倍。需要经过加权求和才能最终获得 4 个特征。

这是因为模型把连续的数值预测问题转换为了离散的分类问题，从 0 到 15 总共 16 个分类。换句话说，模型用 16 个数字来代表范围在 [0, 15] 的一个数字。

例如模型输出的张量维度 [batch, seq, 4 * 16]，转换为 xyhw 的流程如下：

  1. 转换维度到 [batch, seq, 4, 16]
  2. 对 16 这个维度施加 softmax
  3. 对 16 这个维度加权求和，张量维度变换到 [batch, seq, 4]
    • 加权权重为 arange(16)（0、1、2、……、15）

推理时加权相加的做法是用的 1×1 卷积做到的，很聪明的提速方法。

如此可以方便地将问题从连续数值预测转换离散类别分类，进而获得 dfl loss。怪不得要用上 softmax，因为问题是离散的。

顺便一提，模型输出的 4 个特征代表着 x1y1x2y2 距离中心点的偏移量。

使用 CIoU 计算两框重叠率

为了衡量预测框与标注框的差距，进而获得 loss，代码使用了 IoU 的变体 CIoU 进行计算。

关于 IoU 与其变体：

  • IoU（Intersection over Union）通过计算交集面积与并集面积之比得到。缺点是当两框无交集时不能提供“两框距离多远”的信息
    • IoU = Area(Intersection) / Area(Union)
  • GIoU（Generalized IoU），在 IoU 基础上还会找到一个最小的 能够同时包含两框的闭包区域（例如最小外接矩形 C），让 IoU 减去 C 中除并集外的面积占 C 总面积的比例。可见，当一框完全包含另一个框时，GIoU 退化为 IoU
    • GIoU = IoU - ( Area(C) - Area(Union) ) / Area(C)
  • DIoU（Distance IoU），比起 GIoU 使用面积，DIoU 考虑中心点之间的距离。记 Distance 是两个框中心点距离，Diagonal 是包含两个框的最小外接矩形的对角线长度，则公式：
    • DIoU = IoU - ( Distance^2 / Diagonal^2 )
  • CIoU（Complete IoU），在 DIoU 基础上增加一个关于长宽比的惩罚项
    • CIoU = DIoU - alpha * v
    • 其中，
      • v=4π2(atanwgthgt−atanwpredhpred)2v=\frac{4}{\pi^2}(\text{atan}\frac{w_\text{gt}}{h_\text{gt}} - \text{atan}\frac{w_\text{pred}}{h_\text{pred}})^2
      • alpha=v1−IoU+v\text{alpha}=\frac{v}{1-IoU}+v

这样看来 CIoU 性能最好，同时考虑了重叠面积、中心点距离和长宽比三个因素。

YOLO 写的计算 IoU 的代码可以直接拿来用。写得很好，输入两个 tensor 只要保证最后一个维度大小为 4 即可。

from ultralytics.utils.metrics import bbox_iou
avatar
文章
阅读
粉丝