用于PCB缺陷检测系统的YOLOv5算法设计

张振亚 赵成林

（北京邮电大学信息与通信工程学院 北京市 100876）

缺陷目标检测是工业领域很有价值的研究方向。其中印刷电路板（Printed Circuit Board，PCB）的缺陷检测对工业生产十分重要。传统的PCB 缺陷检测方法存在检测速度慢、漏检率高、鲁棒性不佳的种种问题。现代基于CNN 的深度学习目标检测算法对于大目标有很强的检测能力，但对于PCB 中这种小目标缺陷，检测能较弱。小目标不易于检测的原因是，缺陷目标像素过少，在不断卷积的过程中丢失了大部分的有效特征和细节特征。因此，无论是对于传统检测技术还是现代检测技术，小目标检测都是叩待解决的问题。

1 算法设计思路与贡献介绍

为了解决上述问题，对YOLOv5[1]框架进行了研究，提出了ATCSP-YOLOv5 用于PCB 的缺陷检测。实验所使用的PCB 数据集，具有分辨率高，但缺陷目标与图像的相对比例非常小的特点。其中最小的一类缺陷占整个图像的比例仅为0.00893%。结合数据集的特点，ATCSP-YOLOv5 算法采用注意力机制来解决问题。注意力机制在不增加复杂度的前提下，通过给模型分配合适的权重，对小目标的特征信息进行增强。不仅为模型获取了特征图的全局信息，又使目标上下文语义信息得到了丰富。由于图像分辨率较高，所以ATCSP-YOLOv5 选择使用空间注意力机制Self-Attention[2]来进行学习。同时，受到ViT（Vision Transformer）[3]的启发。ATCSP-YOLOv5 借鉴ViT 的思想将Transformer[2]模块嵌入到主干网络中，对主干网络特定深度的特征图做多头自注意力机制的处理。实验证明，此方法有效改善了小目标的检测效果。同时，ATCSP-YOLOv5 还受到CSPNet（跨阶段局部网络）[4]的启发，将SCPNet 的思想与Transformer 模块结合，通过将梯度的变化从头到尾地集成到特征图中，增强了CNN 的学习能力、消除了计算瓶颈、减少了内存成本。其中Transformer 中多头自注意力机制（Multi-Head Attention）和SCPNet 结合的模块为ATCSP（多头自注意力跨阶段局部网络）。

研究所做出的贡献如下：

（1）将多头自注意力机制和SCPNet 的思想结合，生成ATCSP 模块。提高了模型的小目标检测性能。

（2）在特征金字塔部分融合生成感受野更小的检测层，添加了一个小目标探测头；同时去掉了原本感受野最大的探测头，减少了计算量。

（3）通过使用特定的数据增强方法和SRGAN（SRGAN-VGG54）[5]来预处理和扩充数据集，为PCB 数据集的扩充和预处理提供了有效策略。

（4）将训练好的ATCSP-YOLOv5 模型中的PANet[6]换成FPN[7]，对模型进行微调训练，降低了模型的计算量和体积，提升了模型精度。最终将算法封装成检测系统如图1所示。

2 ATCSP-YOLOv5模型

2.1 YOLOv5简介

YOLOv5 有五种不同的模型。考虑到PCB 在流水线检测上实时性和精度的要求，选定YOLOv5s 作为基础框架来进行改进。YOLOv5s 的结构如图2所示。

图2：YOLOv5s 结构图

2.2 ATCSP-YOLOv5模型结构

相比于YOLOv5s，ATCSP-YOLOv5 将主干网络中80×80 特征图大小的C3 结构替换为ATCSP 模块。同时在特征金字塔部分融合生成特征图大小为320×320 的小目标检测层，添加了新的YOLOHead3，删除了YOLOv5s 中的YOLOHead1。ATCSP-YOLOv5 的结构如图3所示。

图3：ATCSP-YOLOv5 结构图

将训练好的ATCSP-YOLOv5 模型中的PANet 替换为FPN 再进行模型参数微调得到ATCSP-YOLOv5_Pruned。

2.3 ATCSP模块

Base Layer 是上一个卷积层的输出。Part1、Part2 是Base Layer 经过特定卷积后分成的两个部分。Part1、Part2一部分不处理，一部分经过原来的模块，然后在最后将两部分进行Concat。ATCSP 的结构如图4所示。

图4：ATCSP 结构图

2.4 Transformer Layer

Transformer Layer 的输入是ATCSP 中的Part2 特征图。将输入的特征图分割成多个16×16 像素的小块。然后将各个小块通过卷积展平成特定维度的一维矢量序列，再给每个矢量加上位置矢量后，便可以送入Transformer Encoder 进行注意力机制的学习了。Transformer Encoder 输出的结果是一个一维矢量需要通过卷积和上采样将其恢复成与输入尺寸相同的特征图。Transformer Layer 的内部结构示意图如图5所示。其中，Transformer Encoder 主要包括两个模块，分别为多头自注意力模块和全连接层（MLP）模块，两模块均使用残差连接。模块中有Norm 层负责归一化处理，有助于网络收敛，防止过拟合。将多头自注意力模块和MLP 作为一个子模块，进行S 次堆叠最终形成Transformer Encoder。

图5：Transformer Layer 示意图

2.5 Multi-Head Attention（多头自注意力机制）

多头自注意力模块是ATCSP 的核心部分。他有效增强了特征图中小目标周围的上下文信息，加强了模型的全局信息关注能力。多头自注意力机制可解释性示意图如图6所示。

图6：多头自注意力可解释性示意图

Multi-Head Attention（多头自注意力）公式如下：

（1）（2）为自注意力机制公式，（3）（4）为多头自注意力机制公式。其中，是特征图生成的输入矩阵，C是特征图的通道数，Din是输入矢量的维度；TQ,TK,TV是可训练的矩阵（本质上是全连接层的权重矩阵）对Xin做线性变换生成V,Q,K作为AT的输入。ATHeadi是多头自注意力编码器中的第i个注意力头部；Dmodel是Transformer Encoder 输入矢量的维度；N是Encoder中的ATHead的数量；

3 实验

3.1 数据集和预处理

训练使用的数据集是来自北京大学智能机器人开放实验室所提供的公共数据集。数据集中包含693 张图像，包括六种缺陷分别为，漏焊、鼠咬、断路、短路、毛刺、余铜。由于数据集数量相对较少，为了增加训练出模型的鲁棒性，实验对数据集进行了预处理。首先使用SRGAN（超分辨率对抗生成网络）对数据集进行超分辨率数据增强，生成细节纹理更加清晰的数据集，增加了小目标的细节特征信息。然后使用复制增强[8]的方法，在图像中复制多份小目标，增加了小目标在图像中出现的次数。最后在此基础上，再使用图像旋转、镜像对称、增加缺陷和背景的对比度、随机亮度调节、HSV 调节几种方式,将数据集扩充到了8316 张。

3.2 实验细节

在实验方面，训练使用的是GPU 型号是NVIDIARTX3090。实验环境使用的是官方给出的YOLOv5 Pytorch版环境。在训练细节上，按训练集和测试集8:2 的比例进行训练。为了防止过拟合，实验采用提前停止训练的策略。经过实验，发现改进后的模型收敛速度很快，模型在56 轮epoch 左右趋近于收敛。训练结束后，将ATCSP-YOLOv5中的PANet 替换为FPN 再进行微调训练。微调训练时，模型在第22 轮epoch 左右收敛，得到了模型体积更小，检测精度更高的ATCSP-YOLOv5_Pruned 模型。

3.3 算法比较

实验选择了两种最为经典的算法与ATCSP-YOLOv5 进行比较，分别是两阶段算法中的Faster R-CNN 和一阶段算法中的SSD。训练曲线对比和检测精度、检测速度、参数量的对比结果，如图7 和表1所示。

表1：各算法mAP@0.5 和检测速度的比较

很明显可以看出ATCSP-YOLOv5 收敛速度明显快于原始YOLOv5，同时mAP@0.5 增加了3.06%。进行FPN 替换后，ATCSP-YOLOv5_Pruned 相较于原始YOLOv5，mAP@0.5增加了3.24%。由图表易见，ATCSP-YOLOv5 的性能远高于Faster R-CNN 和SSDLite_MobileNetV2。

3.4 检测效果展示

如图8所示。

图8：检测效果展示图

4 结束语

ATCSP-YOLOv5 是用于检测与PCB 数据集特点类似的小目标检测算法。首先，该算法借鉴了一些前沿技术的思想，例如Self-Attention、ViT 以及CSPNet 的思想，并将他们巧妙融合起来，作为算法的核心。然后，又提供了许多用于小目标数据集预处理的有效策略。在此基础上，还通过网络替换微调的尝试，得到了一个检测PCB 数据集精度更高的模型，并将此模型做成了一款检测系统。最后，实验表明ATCSP-YOLOv5 算法明显优于已有算法，并保证了检测的实时性。