工业缺陷分类方法、装置、电子设备及存储介质

技术领域

本申请涉及计算机技术领域，特别涉及一种工业缺陷分类方法、装置、电子设备及存储介质。

背景技术

工业制品在现代社会中无处不在，工业缺陷分类，旨在发现各种工业制品的瑕疵，是保障产品质量、维持生产稳定的重要技术之一，近年来，随着工业成像、计算机视觉和深度学习等新的领域技术层出不穷，基于视觉的工业缺陷检测技术得到了很大的发展，也逐渐成为了针对产品外观质检的一种有效的解决方案。

工业场景中常常会面临更换产品线的情况，而不同产品线需要检测的缺陷会有差异，从而使得新的产线没有大量数据可供训练，且再进行大规模标注，则会带来较大的经济和时间成本，针对这种问题，一个有效的解决方法是使用迁移学习，迁移学习是机器学习中解决训练数据不足这一基本问题的重要工具，其中，它可以通过放松训练数据和测试数据必须是独立且同分布的假设，将知识从源域迁移到目标域，迁移学习并不需要对目标域内的模型进行从零开始的训练，就可以显著降低对目标域内训练数据和训练时间的需求，因此这将对许多由于训练数据不足而难以改进的工业缺陷分类任务产生巨大的积极影响。

其中，工业缺陷分类中的迁移学习的学习过程，如图1所示，给定一个基于目标域的缺陷分类任务，通过学习源域的可迁移知识来提高目标域的缺陷分类性能，其中，源域有充足的数据集，所使用的深度学习模型能够收敛，达到较高的性能，而目标域的数据集规模远小于源域数据集规模。此外，工业缺陷分类与常规图像分类任务不同，常规分类样本通常含有大多数类别相关信息，但工业缺陷分类任务常常伴随着检测的部分，通常通过图片局部信息进行缺陷认定，因此，目标域网络模型更易过拟合。

发明内容

本申请提供一种工业缺陷分类方法、装置、电子设备及存储介质，以解决相关技术中工业缺陷分类迁移学习中源域与目标域数据分布差异带来的训练效率低以及已有迁移学习方案精度不高等问题。

本申请第一方面实施例提供一种工业缺陷分类方法，包括以下步骤：获取目标域的训练数据集；设计适用于迁移学习的注意力原型网络；利用训练数据集训练目标域的分类模型，直到满足预设训练条件时，得到训练完成的分类模型，其中，训练阶段，以源域训练所得缺陷分类模型作为目标域训练模型的预训练模型，将预训练模型中预设层的可学习参数进行冻结；对预设层中批量归一化层的统计参数解冻，且增加批量归一化动量；将工业缺陷数据输入所述训练完成的分类模型，输出工业缺陷的实际类别。

可选地，在本申请的一个实施例中，在所述获取目标域的训练数据集之后，还包括：设计在线数据增广模块，训练过程中，在每个批次的数据进入网络之前，经过不同增强程度、不同数据增强手段以增强目标域数据的多样性。

可选地，在本申请的一个实施例中，所述设计适用于迁移学习的注意力原型网络，包括：数据经过第一层卷积、批量归一化、激活函数以及最大池化层后，添加注意力模块，所述注意力模块包括全连接层或卷积层、池化层，经过特征提取下采样层后，添加注意力机制，所述注意力机制采用注意力分数，得到最后的特征向量；将所述特征向量输入模型分类全连接层，采用带边缘的激活函数，公式为：

，

其中，

可选地，在本申请的一个实施例中，所述根据工业迁移学习数据量小，数据分布有差异等问题添加训练策略，包括：训练阶段，以源域训练所得缺陷分类模型作为目标域训练模型的预训练模型，将预训练模型中预设层的可学习参数进行冻结，以引入源域知识；对批量归一化中的统计参数解冻，且增加批量归一化动量。

本申请第二方面实施例提供一种工业缺陷分类模型训练装置，包括：获取模块，用于获取目标域的训练数据集；加载模块，用于加载适用于迁移学习的注意力原型网络；训练模块，用于利用训练数据集训练目标域的分类模型，直到满足预设训练条件时，得到训练完成的分类模型，其中，训练阶段，以源域训练所得缺陷分类模型作为目标域训练模型的预训练模型，将预训练模型中预设层的可学习参数进行冻结；对预设层中批量归一化层的统计参数解冻，且增加批量归一化动量；推理模块，用于将工业缺陷数据输入所述训练完成的分类模型，输出工业缺陷的实际类别。

可选地，在本申请的一个实施例中，工业缺陷分类装置可以对上述方法提到的内容进行设置。

本申请第三方面实施例提供一种电子设备，包括：存储器、处理器及存储在所述存储器上并可在所述处理器上运行的计算机程序，所述处理器执行所述程序，以实现如上述实施例所述的工业缺陷分类方法。

本申请第四方面实施例提供一种计算机可读存储介质，其上存储有计算机程序，该程序被处理器执行，以用于实现如上述实施例所述的工业缺陷分类方法。

由此，本申请至少具有如下有益效果：

本申请实施例可以通过原型设计、有效的特征迁移、数据增强、正则操作，更快的训练出精度更高的模型，从而有效提高训练的效率以及分类结果的精度，能够解决不同工业缺陷分类存在的各种问题。

本申请附加的方面和优点将在下面的描述中部分给出，部分将从下面的描述中变得明显，或通过本申请的实践了解到。

附图说明

本申请上述的和/或附加的方面和优点从下面结合附图对实施例的描述中将变得明显和容易理解，其中：

图1为相关技术中缺陷分类迁移学习的流程图；

图2为根据本申请实施例提供的工业缺陷分类方法的流程图；

图3为根据本申请实施例提供的原型网络结构示意图；

图4为根据本申请实施例提供的新旧产线垫片的示例图；

图5为根据本申请实施例提供的工业缺陷分类装置的方框示意图；

图6为根据本申请实施例提供的电子设备的结构示意图。

具体实施方式

下面详细描述本申请的实施例，实施例的示例在附图中示出，其中自始至终相同或类似的标号表示相同或类似的元件或具有相同或类似功能的元件。下面通过参考附图描述的实施例是示例性的，旨在用于解释本申请，而不能理解为对本申请的限制。

下面参考附图描述本申请实施例的工业缺陷分类方法、装置、电子设备及存储介质。

其中，图2为本申请实施例所提供的一种工业缺陷分类方法的流程示意图。

如图2所示，该工业缺陷分类方法包括以下步骤：

在步骤S101中，获取目标域的训练数据集，源域分类模型（以ResNet50为基础网络结构）。

在步骤S102中，构建适用于迁移学习的注意力原型网络。

在步骤S103中，利用训练数据集训练目标域的分类模型，直到满足预设训练条件时，得到训练完成的分类模型，其中，训练阶段，以源域训练所得缺陷分类模型作为目标域训练模型的预训练模型，将预训练模型中预设层的可学习参数进行冻结。以此引入源域知识；考虑到源域数据分布与目标域数据分布虽然相近，但仍然存在差异，因此对批量归一化中的统计参数解冻，且适量增加批量归一化动量，以到达使模型在训练中，快速学习到目标域数据分布的目的。其中，预设层可以为批次归一化层。

在步骤S104中，将工业缺陷的数据输入训练完成的分类模型，输出工业缺陷的实际类别。

本申请的特点还在于：

步骤S101具体为，获取源域数据集与目标域数据集。设置随机数据增强模块，在训练过程中对每批次的训练数据添加随机数据增强模块，在训练过程中每个批次的数据经过不同的增强程度、不同数据增强手段以增强目标域数据的多样性，相比于训练前对数据集进行数据增广，该方式提高了数据增强方法的灵活性。

步骤S102具体为，图3为本申请实施例设计的原型网络结构示意图，数据经过第一层卷积-批量归一化-激活函数-最大池化（CBRM）池化后，添加注意力模块（CA/SA），经过4个特征提取下采样层后，添加注意力模块，注意力模块包括全连接层或卷积层、池化层，其中注意力机制采用注意力分数Sigmoid，得到最后的特征向量；其中，池化层采用最大池化层。

将特征向量输入模型分类层，其中，分类层由全连接层与激活函数构成。其中，激活函数采用带边缘的Softmax激活函数，提高类间分离程度，减少目标域数据量小引起的过拟合问题。

步骤S103具体按照以下步骤实施：

步骤S103.1、通过源域数据集在基础网络结构训练得到收敛的源域工业缺陷分类模型（简称源域模型）。

步骤S103.2、源域模型作为目标域模型的预训练模型，冻结预设层，但将预设层中批量归一化参数设置为未冻结，并且适量增大批次归一化层更新动量。以此更好地使网络学习目标域数据分布。

步骤S103.3、目标域数据输入模型，将分类层输出结果与训练数据集标签进行损失计算，将损失用于反向训练网络模型。其中损失函数公式为：

，

其中，γ为控制函数曲线的形状，γ的值越大，好分类样本的误差就越小，从而可以把模型的注意力转向那些难以分类的样本上，

，

另外，当γ=0时，损失函数为二分类的交叉熵损失函数，p是模型预测出来的分数，取值在(0,1）；y是真实标签值，取值为1或-1，分别表示正例和负例。

根据本申请实施例提出的工业缺陷分类方法，通过有效的原型网络，特征迁移、数据增强、正则操作，更快的训练出精度更高的模型，从而有效提高训练的效率以及分类结果的精度，能够解决不同工业缺陷分类存在的各种问题。

下面通过一个实际实验案例对本申请提出的工业缺陷分类方法进行阐述。

本申请实施例的实际分类场景实验是在更换产线上，第一条产线部署完成后，在进行下一产线的实施过程中，会出现产线产品虽然相同，但是型号会存在差异，差异主要表现在产品布局，零件形状材料等。

其中，本申请实施例进行实验来充分的证明本申请实施例方法的有效性，具体内容如下：

如图4所示，旧产线和新产线是同一产品的不同型号，其中，旧产线检测项目为圆形垫片是否存在缺陷，新产线增加了矩形垫片是否存在缺陷。

其中，旧产线的数据集为源域数据集，类别两类分别为圆形垫片存在缺陷（NG_C），圆形垫片无缺陷（OK_C），分别有1384、1420张图片数据。

本申请实施例为了对比方法在目标域小数据集情况下的效果，将目标域数据集分为对比训练数据和目标训练数据，数据集数量分别为6223和624，其中有四个类别，除了源域类别，还包含了NG_R和OK_R，分别代表为方垫片有缺陷和方圆垫片无缺陷，目标域数据情况如表1所示。其中，表1为圆、方垫片目标域数据详情表。

表1

训练配置参数批次个数为64，最大迭代数为30。采取3次实验结果平均值来衡量本申请实施例所提出的方法的有效性，结果如表2所示，迁移后的模型训练效率上提升了57.19%，而精度保持在0.5%以内的差异，证明了本申请实施例所提出的方法带来的效率上的提升。其中，表2为圆方垫片迁移学习训练结果表。

表2

注：过检漏检对应的标签顺序为NG_C / NG_R / OK_C / OK_R。

其次参照附图描述根据本申请实施例提出的工业缺陷分类装置。

图5是本申请实施例的工业缺陷分类装置的方框示意图。

如图5所示，该工业缺陷分类装置10包括：获取模块100、加载模块200、训练模块300和推理模块400。

其中，获取模块100，用于获取目标域的训练数据集；加载模块200，用于加载适用于迁移学习的注意力原型网络；训练模块300，用于利用训练数据集训练目标域的分类模型，直到满足预设训练条件时，得到训练完成的分类模型，其中，训练阶段，以源域训练所得缺陷分类模型作为目标域训练模型的预训练模型，将预训练模型中预设层的可学习参数进行冻结；对预设层中批量归一化层的统计参数解冻，且增加批量归一化动量；推理模块400，用于将工业缺陷数据输入训练完成的分类模型，输出工业缺陷的实际类别。

可选地，在本申请的一个实施例中，工业缺陷分类装置10可以对上述方法提到的内容进行设置。

需要说明的是，前述对工业缺陷分类方法实施例的解释说明也适用于该实施例的工业缺陷分类装置，此处不再赘述。

根据本申请实施例提出的工业缺陷分类装置，通过有效的迁移原型网络设计、数据增强、正则操作，更快的训练出精度更高的模型，从而有效提高训练的效率以及分类结果的精度，能够解决不同工业缺陷分类存在的各种问题。

图6为本申请实施例提供的电子设备的结构示意图。该电子设备可以包括：

存储器601、处理器602及存储在存储器601上并可在处理器602上运行的计算机程序。

处理器602执行程序时实现上述实施例中提供的工业缺陷分类方法。

进一步地，电子设备还包括：

通信接口603，用于存储器601和处理器602之间的通信。

存储器601，用于存放可在处理器602上运行的计算机程序。

存储器601可能包含高速RAM（Random Access Memory，随机存取存储器）存储器，也可能还包括非易失性存储器，例如至少一个磁盘存储器。

如果存储器601、处理器602和通信接口603独立实现，则通信接口603、存储器601和处理器602可以通过总线相互连接并完成相互间的通信。总线可以是ISA（IndustryStandard Architecture，工业标准体系结构）总线、PCI（Peripheral Component，外部设备互连）总线或EISA（Extended Industry Standard Architecture，扩展工业标准体系结构）总线等。总线可以分为地址总线、数据总线、控制总线等。为便于表示，图6中仅用一条粗线表示，但并不表示仅有一根总线或一种类型的总线。

可选的，在具体实现上，如果存储器601、处理器602及通信接口603，集成在一块芯片上实现，则存储器601、处理器602及通信接口603可以通过内部接口完成相互间的通信。

处理器602可能是一个CPU（Central Processing Unit，中央处理器），或者是ASIC（Application Specific Integrated Circuit，特定集成电路），或者是被配置成实施本申请实施例的一个或多个集成电路。

本申请实施例还提供一种计算机可读存储介质，其上存储有计算机程序，该程序被处理器执行时实现如上的工业缺陷分类方法。

在本说明书的描述中，参考术语“一个实施例”、“一些实施例”、 “示例”、“具体示例”、或“一些示例”等的描述意指结合该实施例或示例描述的具体特征、结构、材料或者特点包含于本申请的至少一个实施例或示例中。在本说明书中，对上述术语的示意性表述不是必须针对的是相同的实施例或示例。而且，描述的具体特征、结构、材料或者特点可以在任一个或N个实施例或示例中以合适的方式结合。此外，在不相互矛盾的情况下，本领域的技术人员可以将本说明书中描述的不同实施例或示例以及不同实施例或示例的特征进行结合和组合。

此外，术语“第一”、“第二”仅用于描述目的，而不能理解为指示或暗示相对重要性或者隐含指明所指示的技术特征的数量。由此，限定有“第一”、“第二”的特征可以明示或者隐含地包括至少一个该特征。在本申请的描述中，“N个”的含义是至少两个，例如两个，三个等，除非另有明确具体的限定。

流程图中或在此以其他方式描述的任何过程或方法描述可以被理解为，表示包括一个或N个用于实现定制逻辑功能或过程的步骤的可执行指令的代码的模块、片段或部分，并且本申请的优选实施方式的范围包括另外的实现，其中可以不按所示出或讨论的顺序，包括根据所涉及的功能按基本同时的方式或按相反的顺序，来执行功能，这应被本申请的实施例所属技术领域的技术人员所理解。

应当理解，本申请的各部分可以用硬件、软件、固件或它们的组合来实现。在上述实施方式中，N个步骤或方法可以用存储在存储器中且由合适的指令执行系统执行的软件或固件来实现。如，如果用硬件来实现和在另一实施方式中一样，可用本领域公知的下列技术中的任一项或他们的组合来实现：具有用于对数据信号实现逻辑功能的逻辑门电路的离散逻辑电路，具有合适的组合逻辑门电路的专用集成电路，可编程门阵列，现场可编程门阵列等。

本技术领域的普通技术人员可以理解实现上述实施例方法携带的全部或部分步骤是可以通过程序来指令相关的硬件完成，的程序可以存储于一种计算机可读存储介质中，该程序在执行时，包括方法实施例的步骤之一或其组合。

尽管上面已经示出和描述了本申请的实施例，可以理解的是，上述实施例是示例性的，不能理解为对本申请的限制，本领域的普通技术人员在本申请的范围内可以对上述实施例进行变化、修改、替换和变型。