YOLOv6:用于工业应用的单阶段目标检测框架.
YOLOv6 提出了一系列适用于各种工业场景的模型,包括 N/T/S/M/L,考虑到模型的大小,其架构有所不同,以获得更好的精度-速度权衡。本算法专注于检测的精度和推理效率,并在网络结构、训练策略等算法层面进行了多项改进和优化。
简单来说 YOLOv6 开源库的主要特点为:
- 统一设计了更高效的 Backbone 和 Neck:受到硬件感知神经网络设计思想的启发,基于 RepVGG style 设计了可重参数化、更高效的骨干网络 EfficientRep Backbone 和 Rep-PAN Neck。
- 相比于 YOLOX 的 Decoupled Head,进一步优化设计了简洁有效的 Efficient Decoupled Head,在维持精度的同时,降低了一般解耦头带来的额外延时开销。
- 在训练策略上,采用 Anchor-free 的策略,同时辅以 SimOTA 标签分配策略以及 SIoU 边界框回归损失来进一步提高检测精度。
1. 数据增强模块
YOLOv6 目标检测算法中使用的数据增强包括:
- Mosaic 马赛克
- RandomAffine 随机仿射变换
- MixUp
- HSV 颜色空间增强
- 随机水平翻转
在最后 15 个 epoch 的时候将 Mosaic 使用 YOLOv5KeepRatioResize + LetterResize 替代。
2. 网络结构
YOLOv6 N/T/S 模型的网络结构由 EfficientRep + Rep-PAN + Efficient decoupled Head 构成,M/L 模型的网络结构则由 CSPBep + CSPRepPAFPN + Efficient decoupled Head 构成。
YOLOv6 采用了重参数化结构 RepVGG Block 替换掉了原本的 ConvModule,在此基础上,将 CSPLayer 改进为了多个 RepVGG 堆叠的 RepStageBlock(N/T/S 模型)或 BepC3StageBlock(M/L 模型);Head 部分则将回归与分类分支解耦成两个分支进行预测。
(1)backbone
已有研究表明,多分支的网络结构通常比单分支网络性能更加优异,但是这种结构会导致并行度降低进而增加推理延时;相反,单分支网络则具有并行度高、内存占用小的优点,因此推理效率更高。而 RepVGG 则同时具备上述两种结构的优点,在训练时可解耦成多分支拓扑结构提升模型精度,实际部署时可等效融合为单个 3×3 卷积提升推理速度。因此,YOLOv6 基于 RepVGG 重参数化结构设计了高效的骨干网络 EfficientRep 和 CSPBep,其可以充分利用硬件算力,提升模型表征能力的同时降低推理延时。
在 N/T/S 模型中,YOLOv6 使用了 EfficientRep 作为骨干网络,其包含 1 个 Stem Layer 和 4 个 Stage Layer,具体细节如下:
- Stem Layer 中采用 stride=2 的 RepVGGBlock 替换了 stride=2 的 6×6 ConvModule。
- Stage Layer 将 1 个 ConvModule 和 1 个 CSPLayer 分别替换为 1 个 RepVGGBlock 和 1 个 RepStageBlock。其中第一个 RepVGGBlock 会做下采样和 Channel 维度变换,而每个 RepStageBlock 则由 n 个 RepVGGBlock 组成。此外仍然在第 4 个 Stage Layer 最后增加 SPPF 模块后输出。
在 M/L 模型中,由于模型容量进一步增大,直接使用多个 RepVGGBlock 堆叠的 RepStageBlock 结构计算量和参数量呈现指数增长。因此,为了权衡计算负担和模型精度,在 M/L 模型中使用了 CSPBep 骨干网络,其采用 BepC3StageBlock 替换了小模型中的 RepStageBlock 。BepC3StageBlock 由 3 个 1×1 的 ConvModule 和多个子块(每个子块由两个 RepVGGBlock 残差连接)组成。
(2)neck
Neck 部分同样采用 RepStageBlock 或 BepC3StageBlock 对原本的 CSPLayer 进行了替换,需要注意的是,Neck 中 Down Sample 部分仍然使用了 stride=2 的 3×3 ConvModule,而不是像 Backbone 一样替换为 RepVGGBlock。
(3)head
YOLOv6 将分类和回归分支解耦成两个分支进行预测并且去掉了 obj 分支。同时,采用了 hybrid-channel 策略构建了更高效的解耦检测头,将中间 3×3 的 ConvModule 减少为 1 个,在维持精度的同时进一步减少了模型耗费,降低了推理延时。此外,需要说明的是,YOLOv6 在 Backobone 和 Neck 部分使用的激活函数是 ReLU,而在 Head 部分则使用的是 SiLU。
由于 YOLOv6 是解耦输出,分类和 bbox 检测通过不同卷积完成。以 COCO 80 类为例:P5 模型在输入为 640x640 分辨率情况下,其 Head 模块输出的 shape 分别为 $(B,4,80,80), (B,80,80,80), (B,4,40,40), (B,80,40,40), (B,4,20,20), (B,80,20,20)$。
3. 正负样本匹配策略
YOLOv6 采用与 YOLOX 一样的 Anchor-free 无锚范式,省略了聚类和繁琐的 Anchor 超参设定,泛化能力强,解码逻辑简单。在训练的过程中会根据 feature size 去自动生成先验框。网络 bbox 预测的值为 (top, bottom, left, right),解码器将 anchor point 通过四个距离解码到坐标 $(x_1,y_1,x_2,y_2)$。
def decode(points: torch.Tensor, pred_bboxes: torch.Tensor, stride: torch.Tensor) -> torch.Tensor:
"""
将预测值解码转化 bbox 的 xyxy
points (Tensor): 生成的 anchor point [x, y],Shape (B, N, 2) or (N, 2).
pred_bboxes (Tensor): 预测距离四边的距离。(left, top, right, bottom). Shape (B, N, 4) or (N, 4)
stride (Tensor): 特征图下采样倍率.
"""
# 首先将预测值转化为原图尺度
distance = pred_bboxes * stride[None, :, None]
# 根据点以及到四条边距离转为 bbox 的 x1y1x2y2
x1 = points[..., 0] - distance[..., 0]
y1 = points[..., 1] - distance[..., 1]
x2 = points[..., 0] + distance[..., 2]
y2 = points[..., 1] + distance[..., 3]
bboxes = torch.stack([x1, y1, x2, y2], -1)
return bboxes
YOLOv6 采用的标签匹配策略与 TOOD 相同, 前 4 个 epoch 采用 ATSS 作为标签匹配策略的 warm-up , 后续使用 TOOD 算法选择正负样本。
(1)ATSS
ATSS 的匹配策略简单总结为:通过中心点距离先验对样本进行初筛,然后自适应生成 IoU 阈值筛选正样本。 YOLOv6 的实现种主要包括如下三个核心步骤:
- 因为 YOLOv6 是 Anchor-free,所以首先将 anchor point 转化为大小为 5*strdie 的 anchor。
- 对于每一个 GT,在 FPN 的每一个特征层上, 计算与该层所有 anchor 中心点距离(位置先验), 然后优先选取距离 topK 近的样本,作为 初筛样本。
- 对于每一个 GT,计算其 初筛样本 的 IoU 的均值 mean与标准差 std,将 mean + std 作为该 GT 的正样本的 自适应 IoU 阈值 ,大于该 自适应阈值 且中心点在 GT 内部的 anchor 才作为正样本,使得样本能够被 assign 到合适的 FPN 特征层上。
# 1. 首先将anchor points 转化为 anchors
# priors为(point_x,point_y,stride_w,stride_h), shape 为(N,4)
cell_half_size = priors[:, 2:] * 2.5
priors_gen = torch.zeros_like(priors)
priors_gen[:, :2] = priors[:, :2] - cell_half_size
priors_gen[:, 2:] = priors[:, :2] + cell_half_size
priors = priors_gen
# 2. 计算 anchors 与 GT 的 IoU
overlaps = self.iou_calculator(gt_bboxes.reshape([-1, 4]), priors)
# 3. 计算 anchor 与 GT 的中心距离
distances, priors_points = bbox_center_distance(
gt_bboxes.reshape([-1, 4]), priors)
# 4. 根据中心点距离,在 FPN 的每一层选取 TopK 临近的样本作为初筛样本
is_in_candidate, candidate_idxs = self.select_topk_candidates(
distances, num_level_priors, pad_bbox_flag)
# 5. 对于每一个 GT 计算其对应初筛样本的均值与标准差的和, 作为该GT的样本阈值
overlaps_thr_per_gt, iou_candidates = self.threshold_calculator(
is_in_candidate, candidate_idxs, overlaps, num_priors, batch_size,
num_gt)
# 6. 筛选大于阈值的样本作为正样本
is_pos = torch.where(
iou_candidates > overlaps_thr_per_gt.repeat([1, 1, num_priors]),
is_in_candidate, torch.zeros_like(is_in_candidate))
# 6. 保证样本中心点在 GT 内部且不超图像边界
pos_mask = is_pos * is_in_gts * pad_bbox_flag
由于 ATSS 是属于静态标签匹配策略,其选取正样本的策略主要根据 anchor 的位置进行挑选, 并不会随着网络的优化而选取到更好的样本。
(2)TOOD
TOOD 是一种动态样本匹配策略。 在目标检测中,分类和回归的任务最终作用于同一个目标,所以 TOOD 认为样本的选取应该更加关注到对分类以及回归都友好的样本点。
TOOD 的匹配策略简单总结为: 根据分类与回归的分数加权的分数选择正样本。
- 对于每一个 GT,对所有的 预测框 基于 GT类别对应分类分数 与 预测框与 GT 的 IoU 的加权得到一个关联分类以及回归的对齐分数 alignment_metrics。
- 对于每一个 GT,直接基于 alignment_metrics 对齐分数选取 topK 大的作为正样本。
# 1. 基于分类分数与回归的 IoU 计算对齐分数 alignment_metrics
alignment_metrics = bbox_scores.pow(self.alpha) * overlaps.pow(
self.beta)
# 2. 保证中心点在 GT 内部的 mask
is_in_gts = select_candidates_in_gts(priors, gt_bboxes)
# 3. 选取 TopK 大的对齐分数的样本
topk_metric = self.select_topk_candidates(
alignment_metrics * is_in_gts,
topk_mask=pad_bbox_flag.repeat([1, 1, self.topk]).bool())
因为在网络初期参数随机, 分类分数 和 预测框与 GT 的 IoU 都不准确,所以需要经过前 4 个 epoch 的 ATSS 的 warm-up。经过预热之后的 TOOD 标签匹配策略就不使用中心距离的先验, 而是直接对每一个GT 选取 alignment_metrics 中 topK 大的样本作为正样本。
4. Loss函数
参与 Loss 计算的共有两个值:loss_cls 和 loss_bbox。Classes loss使用的是VarifocalLoss;BBox loss对于l/m/s使用的是 GIoULoss, t/n 用的是 SIoULoss。权重比例是:loss_cls : loss_bbox = 1 : 2.5。
(1)分类损失函数 VarifocalLoss
Varifocal Loss是将 预测框与 GT 的 IoU 软化作为分类的标签,使得分类分数关联回归质量, 使其在后处理 NMS 阶段有分类回归一致性很强的分值排序策略,以达到选取优秀预测框的目的。
\[{VFL}(p,q)= \begin{cases} -q(q\log(p) +(1-q)\log(1-p)), & q > 0 \\ -\alpha p^\gamma \log(1-p), & q = 0 \end{cases}\]其中 $q$ 是预测 bboxes 与 GT 的 IoU,使用软标签的形式作为分类的标签。 $p\in[0,1]$ 表示分类分数。
- 对于负样本,即当 $q = 0$ 时,标准交叉熵部分为 $-\log(p)$,负样本权重使用 $\alpha p^\gamma$ 作为 focal weight 使样本聚焦于困难样本上,这与 Focal Loss 基本一致。
- 对于正样本,即当 $q > 0$ 时,首先计算标准二值交叉熵部分 $-(q\log(p) +(1-q)\log(1-p))$, 但是针对正样本的权重设置,Varifocal Loss 认为在网络的学习过程中正样本相对于负样本的学习信号来说更为重要,所以使用了分类的标签 $q$, 即 IoU 作为 focal weight, 使得聚焦到具有高质量的样本上。
YOLOv6 中的 Varifocal Loss 公式采用 TOOD 中的 Task ALignment Learning (TAL), 将预测的 IoU 根据之前标签匹配策略中的分类对齐度 alignment_metrics 进行了归一化, 得到归一化 $\hat{t}$。 具体实现方式为:对于每一个 GT,找到所有样本中与 GT 最大的 IoU,具有最大 alignment_metrics 的样本位置的 $\max(Iou)$:
\[\hat{t} = AlignmentMetrics / max(AlignmentMetrics) * max(IoU)\]YOLOv6 分类损失损失函数为:
\[{VFL}(p,\hat{t})= \begin{cases} -\hat{t}(\hat{t}log(p) +(1-\hat{t})log(1-p)), & \hat{t} > 0 \\ -\alpha p^\gamma log(1-p), & \hat{t} = 0 \end{cases}\](2)回归损失函数 GIoU Loss / SIoU Loss
在 YOLOv6 中,针对不同大小的模型采用了不同的回归损失函数,其中 l/m/s使用的是 GIoULoss, t/n 用的是 SIoU Loss。
由于之前的GIoU, CIoU, DIoU 都没有考虑预测框向 GT 框回归的角度,然而角度也确实是回归中一个重要的影响因素。
SIoU 损失主要由四个度量方面组成:
- IoU成本
- 角度成本
- 距离成本
- 形状成本
角度成本是指图中预测框 $B$ 向 $B^{GT}$ 的回归过程中, 尽可能去使得优化过程中的不确定性因素减少,比如现将图中的角度 $\alpha$ 或者 $\beta$ 变为 $0$ ,再去沿着 $x$ 轴或者 $y$ 轴去回归边界。
在 YOLOv6 中,由于额外的置信度预测头可能与 Aligned Head 有所冲突,经实验验证在不同大小的模型上也都有掉点, 所以最后选择弃用 Objectness 分支。