CN115496738A - 基于Transformer与SE的磁瓦缺陷检测方法 - Google Patents

Abstract

本发明公开了一种基于Transformer与SE的磁瓦缺陷检测方法，包括以下步骤：对磁瓦图像进行数据增强得到数据集，将数据集划分为训练集、测试集；构建磁瓦缺陷检测模型；进行磁瓦缺陷检测模型训练，得到训练好的磁瓦缺陷检测模型；采用测试集进行磁瓦缺陷检测，得到检测结果。本发明在主干网络中增加了Transformer模块与注意力机制SE，Transformer模块充分发掘了深度学习神经网络的特性，利用了自注意力机制实现了快速并行计算，减少了模型的训练时间并提高了磁瓦缺陷的检测精度，注意力机制SE解决了在卷积池化过程中feature map的不同通道所占的重要性不同带来的损失问题。

Description

基于Transformer与SE的磁瓦缺陷检测方法

技术领域

本发明涉及一种基于Transformer与SE的磁瓦缺陷检测方法。

背景技术

铁氧体磁瓦是以铁氧体为主要成分，经高温烧结而制成的永磁材料，是永磁电机产生磁场的关键部件。磁瓦在生产制造过程中，容易出现裂纹、欠磨、偏磨、起级和崩块等多种缺陷。磁瓦质量对永磁电机的性能有较大影响，必须剔除存在表面缺陷的磁瓦。磁瓦磨削加工后的光滑弧形表面和倒角面对光照方向敏感，容易使磁瓦成像背景不均匀；磁瓦色泽灰暗，导致缺陷与磁瓦本身的对比度低，细小缺陷也易受磨削纹理干扰；磁瓦规格和检测面多，形状差异大，检测节拍快，所以磁瓦表面缺陷检测难度很大，国内外对磁瓦表面缺陷检测的研究，距离工程实用化仍有较大差距。

一直以来，磁瓦表面缺陷检测都是以人工目视检测为主，其主要缺点是效率低，成本高，产品检测质量不稳定。目前的多数研究多为根据磁瓦的缺陷对磁瓦进行分类，没有对磁瓦的缺陷位置进行具体的定位。而一些主流的目标检测方法，如YOLOv5方法，没有充分利用图像各个通道的信息，对磁瓦缺陷的目标检测精度不高。

发明内容

为了解决上述技术问题，本发明提供一种算法简单、检测精度高的基于Transformer与SE的磁瓦缺陷检测方法。

本发明解决上述技术问题的技术方案是：一种基于Transformer与SE的磁瓦缺陷检测方法，包括以下步骤：

第一步：对磁瓦图像进行数据增强得到数据集，将数据集划分为训练集、测试集；

第二步：构建磁瓦缺陷检测模型；基于YOLOv5方法进行改进，在其主干网络中增加了Transformer模块与SE注意力机制，从而得到磁瓦缺陷检测模型；

第三步：采用训练集进行磁瓦缺陷检测模型训练，得到训练好的磁瓦缺陷检测模型；

第四步：基于训练好的磁瓦缺陷检测模型，采用测试集进行磁瓦缺陷检测，得到检测结果。

上述基于Transformer与SE的磁瓦缺陷检测方法，所述步骤一中，对图像进行Mosaic数据增强，从训练集中随机取出四张图片进行裁剪，再拼接到一张图上作为训练数据，这样丰富了图片的背景，并且四张图片拼接在一起提高了一次训练所选取的样本数batch size。

上述基于Transformer与SE的磁瓦缺陷检测方法，所述步骤二中，磁瓦缺陷检测模型包括三部分，分别为主干部分Backbone、中间部分Neck、输出部分Output；

Backbone中包含依次连接的第一个Conv、第二个Conv、第一个C3、第三个Conv、第二个C3、第四个Conv、第三个C3、第五个Conv、C3TR、SPPF和SE；Conv包括二维卷积Conv2d、批量标准化BatchNormalization和激活函数HardSwish；C3是带有3个Conv模块的残差结构；SPPF为SPP的快速实现版本；SPP为空间金字塔池化，将任意大小的特征图转换成固定大小的特征向量；C3TR为3个Transformer编码器的残差结构；SE为注意力机制；

Output包含一个深度为256的矩阵的卷积模块Conv256、一个深度为512的矩阵的卷积模块Conv512和一个深度为1024的矩阵的卷积模块Conv1024；

Neck中包含第六个Conv、第七个Conv、第八个Conv、第九个Conv、第四个C3、第五个C3、第六个C3、第七个C3、第一个Concat、第二个Concat、第三个Concat、第四个Concat、第一个UpSample、第二个UpSample，Backbone中SE与第六个Conv连接，第六个Conv与第一个UpSample连接，第一个UpSample分别连接第一个Concat、第二个Concat，第一个Concat经第四个C3后连接Output的Conv1024，Backbone的第三个C3连接第二个Concat，第二个Concat、第五个C3、第七个Conv、第二个UpSample、第三个Concat依次连接，第三个Concat分别连接第六个C3、Backbone的第二个C3，第六个C3与Output的Conv256连接，第六个C3、第八个Conv、第四个Concat、第七个C3、第九个Conv、第一个Concat依次连接，第二个UpSample与第四个Concat连接，第七个C3与Output的Conv512连接；Concat表示将网络进行连接；UpSample为上采样，目的将图片进行放大。

上述基于Transformer与SE的磁瓦缺陷检测方法，所述步骤二中，注意力机制SE的过程为：特征图输入后，通过基于特征图的宽度和高度进行全局平均池化Global Avg Pool的同时对空间特征降维；紧接着使用全连接层Fully Connected和非线性激活函数Non-linear建立通道间的连接；然后经过Sigmoid激活函数获得归一化权重，最后通过乘法逐通道加权到原始特征图的每一个通道上，完成通道注意力对原始特征的重新标定，经过全局平均池化获得全局的感受野，在第一次全连接时通过减少特征图的维度，减少了参数和计算量，之后经过非线性激活函数后再通过一个全连接恢复到原来的通道数，完成了通道间相关性的建立。

上述基于Transformer与SE的磁瓦缺陷检测方法，所述步骤二中，Transformer编码器的结构为：

Transformer分为两大部分，分别为编码器和解码器；编码器由自注意力机制、归一化操作、前向传播部分与残差连接组成；解码器由编码-解码注意力机制、归一化操作、前向传播部分与残差连接组成；输入先进行位置编码嵌入数据的位置信息，然后经过两次编码和两次解码，将得到的输出进行线性化处理再经过Softmax函数激活得到最终输出。

本发明的有益效果在于：本发明在现有的YOLOv5方法的基础上，在其主干网络中增加了Transformer模块与注意力机制SE，Transformer模块充分发掘了深度学习神经网络的特性，利用了自注意力机制实现了快速并行计算，减少了模型的训练时间并提高了磁瓦缺陷的检测精度，注意力机制SE解决了在卷积池化过程中feature map的不同通道所占的重要性不同带来的损失问题。

附图说明

图1为本发明的流程图。

图2为本发明磁瓦图片的Mosaic数据增强示例图。

图3为本发明磁瓦缺陷检测模型的结构图。

图4为Transformer编码器的结构图。

图5为Transformer编码器的原理图。

图6为注意力机制SE的原理图。

具体实施方式

下面结合附图和实施例对本发明作进一步的说明。

如图1所示，一种基于Transformer与SE的磁瓦缺陷检测方法，包括以下步骤：

第一步：对磁瓦图像进行数据增强得到数据集，将数据集划分为训练集、测试集。

对图像进行Mosaic数据增强，从训练集中随机取出四张图片进行裁剪，再拼接到一张图上作为训练数据，这样丰富了图片的背景，并且四张图片拼接在一起变相提高了一次训练所选取的样本数batch size，在进行batch normalization的时候也会计算四张图片。磁瓦图片的Mosaic数据增强示例如图2所示。

第二步：构建磁瓦缺陷检测模型；基于YOLOv5方法进行改进，在其主干网络中增加了Transformer模块与注意力机制，从而得到磁瓦缺陷检测模型。

如图3所示，磁瓦缺陷检测模型包括三部分，分别为主干部分Backbone、中间部分Neck、输出部分Output；

Backbone中包含依次连接的第一个Conv、第二个Conv、第一个C3、第三个Conv、第二个C3、第四个Conv、第三个C3、第五个Conv、C3TR、SPPF和SE；Conv包括二维卷积Conv2d、批量标准化BatchNormalization和激活函数HardSwish；C3是带有3个Conv模块的残差结构；SPPF为SPP的快速实现版本；SPP为空间金字塔池化，将任意大小的特征图转换成固定大小的特征向量；C3TR为3个Transformer编码器的残差结构；SE为注意力机制；

Output包含一个深度为256的矩阵的卷积模块Conv256、一个深度为512的矩阵的卷积模块Conv512和一个深度为1024的矩阵的卷积模块Conv1024；

Neck中包含第六个Conv、第七个Conv、第八个Conv、第九个Conv、第四个C3、第五个C3、第六个C3、第七个C3、第一个Concat、第二个Concat、第三个Concat、第四个Concat、第一个UpSample、第二个UpSample，Backbone中SE与第六个Conv连接，第六个Conv与第一个UpSample连接，第一个UpSample分别连接第一个Concat、第二个Concat，第一个Concat经第四个C3后连接Output的Conv1024，Backbone的第三个C3连接第二个Concat，第二个Concat、第五个C3、第七个Conv、第二个UpSample、第三个Concat依次连接，第三个Concat分别连接第六个C3、Backbone的第二个C3，第六个C3与Output的Conv256连接，第六个C3、第八个Conv、第四个Concat、第七个C3、第九个Conv、第一个Concat依次连接，第二个UpSample与第四个Concat连接，第七个C3与Output的Conv512连接；Concat表示将网络进行连接；UpSample为上采样，目的将图片进行放大。

如图4、图5所示，Transformer编码器的结构为：

Transformer分为两大部分，分别为编码器和解码器；编码器由自注意力机制、归一化操作、前向传播部分与残差连接组成；解码器由编码-解码注意力机制、归一化操作、前向传播部分与残差连接组成；输入先进行位置编码嵌入数据的位置信息，然后经过两次编码和两次解码，将得到的输出进行线性化处理再经过Softmax函数激活得到最终输出。

x₁,x₂,x₃,x₄是一个输入序列，经过自注意力机制计算后得到y₁,y₂,y₃,y₄也是一个序列；w₂₁,w₂₂,w₂₃,w₂₄是要更新的权重。Add&Normalize表示将结果相加再对结果进行分布归一化，转换成均值为0方差为1的正态分布。Feed Forward为前向传播；Softmax为激活函数，其主要用于分类任务，往往放在最后一层，其会生成一个N维向量，数据集中的每一类对应一个向量(N一般为数据类中数据的类别数目)，Linear为线性化。

如图6所示，SE表示注意力机制，其目的是获得在通道维度上的注意力，注意力机制SE的过程为：特征图输入后，通过基于特征图的宽度和高度进行全局平均池化GlobalAvg Pool的同时对空间特征降维；紧接着使用全连接层Fully Connected和非线性激活函数Non-linear建立通道间的连接；然后经过Sigmoid激活函数获得归一化权重，最后通过乘法逐通道加权到原始特征图的每一个通道上，完成通道注意力对原始特征的重新标定，经过全局平均池化获得全局的感受野，在第一次全连接时通过减少特征图的维度，减少了参数和计算量，之后经过非线性激活函数后再通过一个全连接恢复到原来的通道数，完成了通道间相关性的建立。

第三步：采用训练集进行磁瓦缺陷检测模型训练，得到训练好的磁瓦缺陷检测模型。

第四步：基于训练好的磁瓦缺陷检测模型，采用测试集进行磁瓦缺陷检测，得到检测结果。

本发明通过通过python语言编写，使用Pytorch框架。对磁瓦缺陷数据集进行测试。测试结果如下表所示：

表1实验对比

根据表1可知：使用Transformer模块对YOLOv5进行改进后，精确率相对于只使用YOLOv5方法提高了5.7％，mAP@0.5提高了2.4％，但是召回率有所下降。添加了SE注意力机制后，精确率进一步提升了0.8％，mAP@0.5进一步提升了2.5％，虽然召回率有所提升，还是低于原方法。mAP@0.5为衡量在目标检测的重要指标，综上所述，在提升了精确率和mAP@0.5的情况下，召回率的下降在可接收范围内。

Claims (5)

1.一种基于Transformer与SE的磁瓦缺陷检测方法，其特征在于，包括以下步骤：

第一步：对磁瓦图像进行数据增强得到数据集，将数据集划分为训练集、测试集；

第二步：构建磁瓦缺陷检测模型；基于YOLOv5方法进行改进，在其主干网络中增加了Transformer模块与注意力机制，从而得到磁瓦缺陷检测模型；

第三步：采用训练集进行磁瓦缺陷检测模型训练，得到训练好的磁瓦缺陷检测模型；

第四步：基于训练好的磁瓦缺陷检测模型，采用测试集进行磁瓦缺陷检测，得到检测结果。

2.根据权利要求1所述的基于Transformer与SE的磁瓦缺陷检测方法，其特征在于，所述步骤一中，对图像进行 Mosaic 数据增强，从训练集中随机取出四张图片进行裁剪，再拼接到一张图上作为训练数据，这样丰富了图片的背景，并且四张图片拼接在一起提高了一次训练所选取的样本数batch size。

3.根据权利要求2所述的基于Transformer与SE的磁瓦缺陷检测方法，其特征在于，所述步骤二中，磁瓦缺陷检测模型包括三部分，分别为主干部分Backbone、中间部分Neck、输出部分Output；

Backbone中包含依次连接的第一个Conv、第二个Conv、第一个C3、第三个Conv、第二个C3、第四个Conv、第三个C3、第五个Conv、C3TR、SPPF和SE；Conv包括二维卷积Conv2d、批量标准化BatchNormalization和激活函数HardSwish；C3是带有3个Conv模块的残差结构；SPPF为SPP的快速实现版本；SPP为空间金字塔池化，将任意大小的特征图转换成固定大小的特征向量；C3TR为3个Transformer编码器的残差结构；SE为注意力机制；

Output包含一个深度为256的矩阵的卷积模块Conv256、一个深度为512的矩阵的卷积模块Conv512和一个深度为1024的矩阵的卷积模块Conv1024；

Neck中包含第六个Conv、第七个Conv、第八个Conv、第九个Conv、第四个C3、第五个C3、第六个C3、第七个C3、第一个Concat、第二个Concat、第三个Concat、第四个Concat、第一个UpSample、第二个UpSample，Backbone中SE与第六个Conv连接，第六个Conv与第一个UpSample连接，第一个UpSample分别连接第一个Concat、第二个Concat，第一个Concat经第四个C3后连接Output的Conv1024，Backbone的第三个C3连接第二个Concat，第二个Concat、第五个C3、第七个Conv、第二个UpSample、第三个Concat依次连接，第三个Concat分别连接第六个C3、Backbone的第二个C3，第六个C3与Output的Conv256连接，第六个C3、第八个Conv、第四个Concat、第七个C3、第九个Conv、第一个Concat依次连接，第二个UpSample与第四个Concat连接，第七个C3与Output的Conv512连接；Concat表示将网络进行连接；UpSample为上采样，目的将图片进行放大。

4.根据权利要求3所述的基于Transformer与SE的磁瓦缺陷检测方法，其特征在于，所述步骤二中，注意力机制SE的过程为：特征图输入后，通过基于特征图的宽度和高度进行全局平均池化Global Avg Pool的同时对空间特征降维；紧接着使用全连接层FullyConnected和非线性激活函数Non-linear建立通道间的连接；然后经过 Sigmoid 激活函数获得归一化权重，最后通过乘法逐通道加权到原始特征图的每一个通道上，完成通道注意力对原始特征的重新标定，经过全局平均池化获得全局的感受野，在第一次全连接时通过减少特征图的维度，减少了参数和计算量，之后经过非线性激活函数后再通过一个全连接恢复到原来的通道数，完成了通道间相关性的建立。

5.根据权利要求4所述的基于Transformer与SE的磁瓦缺陷检测方法，其特征在于，所述步骤二中，Transformer的结构为：

Transformer分为两大部分，分别为编码器和解码器；编码器由自注意力机制、归一化操作、前向传播部分与残差连接组成；解码器由编码-解码注意力机制、归一化操作、前向传播部分与残差连接组成；输入先进行位置编码嵌入数据的位置信息，然后经过两次编码和两次解码，将得到的输出进行线性化处理再经过Softmax函数激活得到最终输出。

Images