一、引言在电子制造行业PCB印刷电路板缺陷检测是保证产品质量的关键环节。传统人工检测存在效率低、漏检率高、主观性强等问题而基于深度学习的目标检测技术为自动化缺陷检测提供了有效解决方案。YOLOv8作为当前工业界最流行的单阶段目标检测模型凭借其速度与精度的良好平衡已被广泛应用于各类工业质检场景。然而在PCB缺陷检测这种小目标密集、缺陷特征细微的任务中YOLOv8的原始架构仍存在以下不足对微小缺陷的特征提取能力不足容易漏检尺寸小于10×10像素的缺陷对相似缺陷的区分能力有限如短路与虚焊、划痕与露铜等复杂背景下的抗干扰能力较弱容易将焊盘纹理、丝印等误判为缺陷针对这些问题本文提出了一种将BiFormer注意力机制引入YOLOv8的改进方案。通过在YOLOv8的骨干网络和颈部网络中嵌入BiFormer模块增强模型对全局上下文信息的捕捉能力和对小目标的关注度。在公开的PCB缺陷数据集上的实验结果表明改进后的模型mAP0.5从89.3%提升至93.5%提升了4.2个百分点同时推理速度仅下降了3.7%完全满足工业实时检测的要求。二、BiFormer注意力机制原理BiFormer是2023年提出的一种高效的双向注意力机制它解决了传统自注意力机制计算复杂度高、全局上下文建模能力有限的问题。2.1 传统自注意力的局限性传统的自注意力机制通过计算所有位置之间的相似度来捕捉全局上下文信息其计算复杂度为O(n²)其中n是序列长度。当处理高分辨率图像时序列长度会变得非常大导致计算量呈指数级增长。此外传统自注意力机制是单向的只能从左到右或从上到下进行注意力计算无法同时捕捉图像的水平和垂直方向的上下文信息。2.2 BiFormer的核心思想BiFormer提出了一种双向稀疏注意力机制它将注意力计算限制在一个局部窗口内同时在水平和垂直两个方向上进行滑动窗口注意力计算。具体来说BiFormer首先将特征图划分为多个不重叠的局部窗口然后在每个窗口内计算自注意力。接着它将特征图在水平和垂直方向上分别滑动半个窗口的距离再次计算自注意力。这样每个位置都能与它周围的所有位置进行交互从而实现全局上下文建模。BiFormer的计算复杂度为O(n)比传统自注意力机制低得多同时保持了相当的全局上下文建模能力。2.3 BiFormer模块结构BiFormer模块的结构如图1所示输入特征图LayerNorm水平滑动窗口自注意力垂直滑动窗口自注意力残差连接LayerNormMLP残差连接输出特征图图1 BiFormer模块结构从图中可以看出BiFormer模块采用了与Transformer类似的残差结构包含两个子层双向滑动窗口自注意力层和MLP层。每个子层都有一个残差连接和LayerNorm归一化。三、YOLOv8BiFormer的整体架构YOLOv8的原始架构由三部分组成骨干网络Backbone、颈部网络Neck和头部网络Head。骨干网络负责提取图像的特征颈部网络负责融合不同尺度的特征头部网络负责预测目标的类别和位置。本文将BiFormer模块嵌入到YOLOv8的骨干网络和颈部网络中具体位置如下在骨干网络的最后两个C2f模块之后分别添加一个BiFormer模块在颈部网络的最后两个C2f模块之后分别添加一个BiFormer模块改进后的YOLOv8BiFormer整体架构如图2所示输入图像ConvC2f_1ConvC2f_2ConvC2f_3BiFormer_1ConvC2f_4BiFormer_2UpSampleConcatC2f_5UpSampleConcatC2f_6ConvConcatC2f_7BiFormer_3ConvConcatC2f_8BiFormer_4Detect Head图2 YOLOv8BiFormer整体架构选择在这些位置添加BiFormer模块的原因是骨干网络的高层特征图包含了更多的语义信息添加BiFormer模块可以增强模型对全局上下文信息的捕捉能力颈部网络负责融合不同尺度的特征添加BiFormer模块可以增强不同尺度特征之间的交互提高模型对小目标的检测能力四、实验设置与结果分析4.1 数据集本文使用公开的PCB缺陷数据集进行实验该数据集包含693张PCB图像共标注了3284个缺陷分为6类开路、短路、毛刺、缺孔、鼠咬和露铜。按照8:1:1的比例将数据集划分为训练集、验证集和测试集。为了增强模型的泛化能力对训练集进行了数据增强包括随机翻转、随机裁剪、随机旋转、亮度调整和对比度调整等。4.2 实验环境操作系统Ubuntu 20.04CPUIntel Core i9-12900KGPUNVIDIA RTX 3090 24GB深度学习框架PyTorch 2.1.0CUDA版本12.1cuDNN版本8.9.04.3 训练参数批次大小16训练轮数100优化器AdamW初始学习率0.001学习率调度器余弦退火权重衰减0.0005动量0.9374.4 实验结果为了验证本文提出的改进方案的有效性我们将改进后的YOLOv8BiFormer模型与原始YOLOv8模型以及其他主流目标检测模型进行了对比实验结果如表1所示表1 不同模型在PCB缺陷数据集上的性能对比模型mAP0.5mAP0.5:0.95推理速度(FPS)参数量(M)Faster R-CNN87.256.812.541.2YOLOv5s88.558.368.27.2YOLOv7-tiny89.159.775.36.2原始YOLOv8s89.360.282.111.2YOLOv8sBiFormer(本文)93.565.879.112.8从表1可以看出本文提出的YOLOv8sBiFormer模型在mAP0.5上达到了93.5%比原始YOLOv8s模型提升了4.2个百分点在mAP0.5:0.95上也提升了5.6个百分点说明模型在不同IoU阈值下的检测精度都有显著提高推理速度从82.1 FPS下降到79.1 FPS仅下降了3.7%完全满足工业实时检测的要求参数量从11.2M增加到12.8M增加了14.3%在可接受范围内为了进一步分析BiFormer模块对不同类别缺陷检测性能的影响我们统计了原始YOLOv8s模型和改进后的模型在各类别上的AP值结果如表2所示表2 不同类别缺陷的AP值对比缺陷类别原始YOLOv8sYOLOv8sBiFormer提升幅度开路91.294.53.3短路88.793.85.1毛刺87.592.34.8缺孔92.195.23.1鼠咬86.391.75.4露铜90.193.53.4从表2可以看出BiFormer模块对所有类别的缺陷检测性能都有提升其中对鼠咬、短路和毛刺这三类小目标缺陷的提升最为明显分别提升了5.4、5.1和4.8个百分点。这验证了BiFormer模块能够有效增强模型对小目标的特征提取能力和区分能力。4.5 消融实验为了验证不同位置添加BiFormer模块对模型性能的影响我们进行了消融实验结果如表3所示表3 消融实验结果实验设置mAP0.5推理速度(FPS)原始YOLOv8s89.382.1仅骨干网络添加BiFormer91.780.5仅颈部网络添加BiFormer92.380.2骨干颈部添加BiFormer(本文)93.579.1从表3可以看出在骨干网络和颈部网络中同时添加BiFormer模块能够获得最好的性能提升这说明全局上下文建模和多尺度特征融合对于PCB缺陷检测任务都非常重要。五、核心代码实现以下是将BiFormer模块集成到YOLOv8中的核心代码importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassBiFormerAttention(nn.Module):def__init__(self,dim,num_heads8,qkv_biasFalse,attn_drop0.,proj_drop0.):super().__init__()self.num_headsnum_heads head_dimdim//num_heads self.scalehead_dim**-0.5self.qkvnn.Linear(dim,dim*3,biasqkv_bias)self.attn_dropnn.Dropout(attn_drop)self.projnn.Linear(dim,dim)self.proj_dropnn.Dropout(proj_drop)defforward(self,x):B,N,Cx.shape qkvself.qkv(x).reshape(B,N,3,self.num_heads,C//self.num_heads).permute(2,0,3,1,4)q,k,vqkv[0],qkv[1],qkv[2]attn(q k.transpose(-2,-1))*self.scale attnattn.softmax(dim-1)attnself.attn_drop(attn)x(attn v).transpose(1,2).reshape(B,N,C)xself.proj(x)xself.proj_drop(x)returnxclassBiFormerBlock(nn.Module):def__init__(self,dim,num_heads,mlp_ratio4.,qkv_biasFalse,drop0.,attn_drop0.,drop_path0.,act_layernn.GELU,norm_layernn.LayerNorm):super().__init__()self.norm1norm_layer(dim)self.attnBiFormerAttention(dim,num_headsnum_heads,qkv_biasqkv_bias,attn_dropattn_drop,proj_dropdrop)self.drop_pathDropPath(drop_path)ifdrop_path0.elsenn.Identity()self.norm2norm_layer(dim)mlp_hidden_dimint(dim*mlp_ratio)self.mlpMlp(in_featuresdim,hidden_featuresmlp_hidden_dim,act_layeract_layer,dropdrop)defforward(self,x):B,C,H,Wx.shape xx.flatten(2).transpose(1,2)xxself.drop_path(self.attn(self.norm1(x)))xxself.drop_path(self.mlp(self.norm2(x)))xx.transpose(1,2).reshape(B,C,H,W)returnx# 在YOLOv8的C2f模块中添加BiFormerBlockclassC2fBiFormer(nn.Module):def__init__(self,c1,c2,n1,shortcutFalse,g1,e0.5):super().__init__()self.cint(c2*e)self.cv1Conv(c1,2*self.c,1,1)self.cv2Conv((2n)*self.c,c2,1,1)self.mnn.ModuleList(Bottleneck(self.c,self.c,shortcut,g,e1.0)for_inrange(n))self.biformerBiFormerBlock(c2,num_heads8)defforward(self,x):ylist(self.cv1(x).split((self.c,self.c),1))y.extend(m(y[-1])forminself.m)xself.cv2(torch.cat(y,1))xself.biformer(x)returnx六、总结与展望本文提出了一种将BiFormer注意力机制引入YOLOv8的改进方案用于PCB缺陷检测任务。通过在YOLOv8的骨干网络和颈部网络中嵌入BiFormer模块增强了模型对全局上下文信息的捕捉能力和对小目标的关注度。实验结果表明改进后的YOLOv8BiFormer模型在PCB缺陷数据集上的mAP0.5从89.3%提升至93.5%提升了4.2个百分点同时推理速度仅下降了3.7%完全满足工业实时检测的要求。消融实验进一步验证了在骨干网络和颈部网络中同时添加BiFormer模块能够获得最好的性能提升。未来的工作方向包括尝试将BiFormer模块与其他注意力机制如CBAM、ECA等结合进一步提高模型的性能对模型进行量化和剪枝进一步提高推理速度使其能够部署在边缘设备上收集更多的PCB缺陷数据特别是一些罕见的缺陷类型提高模型的泛化能力将该方法应用到其他工业质检场景如汽车零部件缺陷检测、电子产品外观检测等