基于Self-Attention的单样本ConSinGAN模型的工业缺陷样本图像生成

胡文杰，吴晓波，李波*，徐天伦，姚为

（1中南民族大学 计算机科学学院，武汉 430074；2中国外运股份有限公司，北京 100029）

在工业生产中，基于机器视觉的产品质量检测一般是采用有监督训练来建模.深度学习需要充足的样本集用于训练模型，但随着企业生产工艺水平的提高，缺陷产品出现的概率往往很低，因此缺陷样本集一般是小样本.单样本或小样本的缺陷图像对于目标检测而言存在模型训练过拟合的问题.而采用传统的图像增广方法，如仿射变换与裁剪翻转都无法解决图像卷积过程中的平移不变性带来的过拟合情况.鉴于此，本文的研究目标是使用单张工业缺陷样本图像训练生成式对抗网络，以获取能稳定生成指定缺陷的工业缺陷图像生成模型.

在传统单样本图像生成方法中，如基于图像摘要的双向相似性度量［1］等技术，该技术需要大量的先验知识和函数设计.有些方法局限性较高，只对纹理图像进行训练［2］.目前只有少数模型能够在单个“自然”图像上进行训练［3-5］，或仅用于研究小样本图像之间的转换迁移［6］.

从近几年GAN生成样本的工作上来看，单个图像上训练GAN模型的方法仍然相对较少.SinGAN［7］是首个能在无监督条件下将单张自然图像作为训练集进行图像生成训练的对抗生成网络模型.SinGAN模型仅对当前最近阶段进行参数训练，并冻结所有先前阶段的参数.该方法在每个批次的训练阶段是没有不同尺度的特征融合操作，导致网络需要保存所有阶段的特征图参数.

ConSinGAN［8］网络包含一种新的渐进式生长技术，可以同时训练多个阶段，但不是全部阶段，同时在较低阶段使用渐进式较小的学习率.由于单张图像样本同时训练模型的多个阶段，因此该模型称为并发单张图片GAN（Concurrent Single Image GAN）.

ConSinGAN模型使用相邻两阶段的特征图融合与尺度非线性缩放等手段，改进了SinGAN模型的训练时间过高与迭代消耗内存过高的问题.工业生产图片背景较为单一、特征较弱、缺陷与背景间对比度较低.上述特异性样本会导致风景图像生成模型ConSinGAN训练极不稳定.

本文在ConSinGAN模型的基础上进行改进，引入了自注意力机制和结构相似度损失函数，使得模型训练得到的图像不仅评估指标有较好提升，而且训练后的图像质量效果也有较大提高.本文改进工作主要解决了背景纹理特征抽取和缺陷区域特征学习困难这两个问题，主要贡献点有：

（1）引入结构相似度损失函数的混合损失，显著降低了由于纹理缺陷图像背景过于单一，训练容易崩溃的几率（见1.2）.

（2）引入注意力模块（见2.5）提高了工业样本生成的图像性能，极大增加了样本图像生成中缺陷出现的概率，并有效解决了缺陷边缘模糊和背景不一致问题.

1 ConSinGAN模型及本文改进

对于小样本图像来说，为了解决样本空间过小导致的过拟合问题，一般常采用迭代式的网络结构［9］设计来增强样本的特征表达丰富程度.

1.1 ConSinGAN模型设计

ConSinGAN模型设计采用多阶段叠加的策略来优化多尺度特征图的融合.原始图像从一个较小的分辨率开始，进行多次迭代，学习从随机噪声矢量z到低分辨率图像的映射（参见图1中的“生成器：阶段0”）.

图1 ConSinGAN训练过程图Fig.1 The Training processing of ConSinGAN

一旦阶段n的训练收敛，ConSinGAN模型通过添加三个额外的卷积层来增加生成器的大小.该操作有助于训练出一定差异的生成器.每个阶段都从前一阶段获取原始特征作为输入，前一层参数不再被冻结.模型将重复这个过程N（常规尺寸的图片设置N为6）直到达到满足条件的输出分辨率.为了提高生成器的稳定和多样性，在每个阶段向特征添加额外的噪声［10-11］.

复合损失函数设计为对抗损失和重建损失之和，如公式（1）：

其中Ladv（Gn，Dn）是WGAN-GP对抗损失函数［1］，在这里还使用了αLrec（Gn）来重建生成器损失函数，α为缩放系数.对于重建损失函数来说，生成器讲原始图像x0的第n阶段下采样xn作为输入，重建损失函数见公式（2）：

1.2 本文模型的改进

工业生产图片中背景特异性较强，模型训练需要对判别器图形的纹理背景相似度进行更为精确的考量.为了达到更好的纹理学习效果，本文创新性地引入了结构相似度损失函数LSSIM函数［12］公式（3）和（4）到GAN模型中，

其中，G，F为生成器产生的图片和同尺寸缩小的原图，μG，μF为G，F的平均值为G，F的方差，σGF是协方差.c1=（k1L）2，c2=（k2L）2是维持稳定的常数，L是图像的动态范围.k1=0.01，k2=0.03.

改进后模型的复合Loss见公式（5），其中α和β为权重缩放系数.

2 CBAM自注意力模块模型融合

2.1 自注意力机制

众所周知，注意力在人类视觉感知中起着重要作用［13-14］.人类视觉系统的一个重要特性是聚焦于重点部分，而不是整个场景.最近的一些相关研究工作［15-16］将注意力处理和神经网络结合起来，以提高CNN在大规模分类任务中的性能.通过注意力机制添加的特征图关联映射，使得修改后的卷积网络不仅性能良好，而且对噪声输入具有鲁棒性.

本文选用CBAM［17］模块，该模块采用分解后的串行注意力机制，而不是直接计算三维注意力特征图［15-16］.串行计算可以减少特征图的通道数和获得较小的特征图尺寸，因此具有更少的计算和参数开销，设置好输入和输入维度后可以用作现有基础CNN架构的即插即用模块.

2.2 CBAM模块

CBAM模块特征图采用最大池化及平均池化两者的共同组合以获取强度更高的自注意力特征，并同时利用了空间和通道注意.并通过实证验证了利用两者结合的方式优于仅使用通道注意力［18］.注意力过程可概括为公式（6）和（7），

这里⊗表示像素级乘积.在乘积中注意力的数值得到广播，即通道注意力被广播到频域空间.F″是最终的输出特征.图2描绘了每个注意图的计算过程.下面将描述每个注意模块的详细信息.

图2 CBAM模块模型结构图Fig.2 The Structure of CBAM Module

2.3 通道注意力机制

通道注意力作用于输入图像各位置的感知强度，即更关心图像中的缺陷或者前景区域.

CBAM模块使用平均池化和最大池化操作来聚合特征地图的空间信息，生成两个不同的空间上下文特征映射分别表示平均池特征和最大池特征.两者通过维度拼接得到通道注意图Mc∈Rc×1×1.共享网络由多层感知器（MLP）和一个隐藏层组成.为了减少参数开销，将隐藏激活大小设置为Rc/r×1×1，其中r为缩减率.通道注意力的计算如式（8）：

W0∈Rc/r×c且W1∈Rc×c/r.需要注意的是，MLP权重W0和W1为同一个特征图输入，σ表示激活函数.

2.4 空间注意模块

与通道注意不同的是，空间注意侧重于“何处”是信息部分，该模块与通道注意获取的注意力特征是互补的.空间注意力模块首先沿特征通道应用平均池和最大池操作，并将它们连接起来以生成有效的特征描述符.

卷积层生成空间注意图Ms（F）∈RH×W表示加强或抑制的位置.该模块通过两个池化操作聚合特征映射的通道信息，生成两个二维映射分别表示通道中的平均池特征和最大池特征.然后，通过卷积融合生成CBAM的二维空间注意力图.空间注意力的计算如式（9）：

其中σ表示激活函数，f7×7表示滤波器大小为7×7的卷积运算.

2.5 本文提出的CBAM模型融合设计

一般注意力机制常被引入到目标检测及分割［19-21］等研究中，本文成功引入注意力机制到单样本生成模型ConSinGAN中，并在3.3节中通过实验详细探究了CBAM模块放置位置得到的模型效果.

具体的CBAM模块结构如图3所示.作为一个即插即用的注意力模块，理论上只需要对齐通道数即可将多通道注意力融入模型中.

图3 可放置到指定位置的CBAM模块Fig.3 CBAM modules that can be placed in designated location

鉴于单张图片迭代式多尺度的作为判别器的采样图，需要将CBAM模块放入不同级别的混合噪声向量初始化判别器后.而在生成器中，为了提高模型的生成能力，防止过拟合的导致缺陷表现单一，CBAM模块不能放入生成器内.

表1详述了判别器与生成器神经网络设计.生成器的输入为混合随机向量，然后通过全卷积操作（Conv）和正则处理（Norm）最后加上一个正切激活函数来保持出的平滑性.判别器（Discriminator，D）的输入为生成器的输出，在卷积运算中插入CBAM模块，提高了判别器对于特定区域（缺陷）的判别能力.可以反向提高生成器产生缺陷区域的效果.

表1 生成器和判别器神经网络结构表Tab.1 G and D structure in neural network

3 实验与结果分析

3.1 MVTec AD公开数据集

本文实验基于MVTec AD公开数据集［22］开展实验，该数据集包含5354张不同对象和纹理类别的高分辨率彩色图像.它包含用于训练的正常图像，即无缺陷的图像，和用于测试的缺陷图像.缺陷图像包含在超过70种不同类型的缺陷，如划痕，凹痕，污染和各种结构变化.具体样本统计分布情况见表2，其中训练集为无缺陷正常图像，测试集包含各类缺陷样本.

表2 MVTec AD数据集样本种类分布情况表Tab.2 Distribution of sample types in MVtec

具体来说，图像分为两类，一类是物体表面，包含原木板表面、麻布块等，另一类是工业单物件产品，例如含有漆包线、半导体元件等.每种产品含有该产物在生产环节中产生的常见缺陷样本，并且缺陷种类较多.

图4展示了较为显著的缺陷样本，图4上部为三角管（transistor），包含的缺陷种类有缺脚、开胶、单脚缺焊和脱焊.图4下部为原木板表面（woods），（a）到（e）分别表示正常（normal）、穿孔（hole）、染色（color）、水渍（liquid）与混合缺陷（combined）.

图4 MVTec AD中三角管和木板的缺陷种类样图Fig.4 Sample of defect types of translator and wood in MVTec AD

3.2 实验环境与参数设置

实验的环境主要包括：硬件环境是ubuntu18.04操作系统，NVIDIA Tesla P40 24 GB显卡，内存大小为32GB，CPU为Intel X5；软件环境为python 3.9下pytorch1.9和cuda10.2.实验设定阶段数量为6.

对于不同阶段来说，中间生成的图像的尺度差异是较大，因此采用学习率缩放是十分必要的.在早期阶段的训练中，使用小的学习率有助于减少过拟合，即早期大的学习率会导致模型迅速崩溃.所以使用一个在学习率η上需要加上一个缩放因子δ.对于阶段n的生成器Gn来说，学习率为δnη.阶段n以学习率δ0η进行训练，阶段n-2以学习率δ1η进行培训，阶段n-2使用δ2η等，在实验中设置δ=0.1在图像保真度和多样性之间取得了良好的权衡.

ConSinGAN网络共有6个阶段的迭代，每个阶段的迭代共包含3个全连接的卷积层和相应的激活和池化层.但是每个阶段对应的空间特征图的尺度都有等比例的缩放.由多尺度缩放系数要求可求得参数量约在66万个.

在引入自注意力机制后，双通道注意力只经过平均池化和最大池化操作后拼接在一起.并且为了和每个阶段的特征图尺度对应，增加的模块参数量约为原ConSinGAN模型的18.2%，达到了78万个.引入SSIM结构相似度Loss后，模型的结构没有变化，损失函数的计算相较于原函数只有常数时间的提高，因此对模型的计算量的增加可以忽略，训练时间最多延长了3%.

3.3 CBAM模块模型融合能力探究实验及分析

CBAM模块是一个对齐特征图输入输出即可成功引入的自注意力模块，可以放入判别器网络结构中.该模块其具有一定的灵活性，可以放入判别器的头部、中间或尾部.

为了获取该模块的最佳性能表现，将对该模块分别放置在判别器的头中尾三个部分分别进行对比实验.具体的实验结果见表3.

表3 CBAM插入位置对应缺陷生成情况表Tab.3 Defect generation corresponding to CBAM placement position

训练成功率表现了模型的训练稳定性，鉴于输入图像为单样本，所以训练失败的最主要原因是损失函数值过小导致的梯度消失问题，而自注意力机制的成功引入增强了模型对工业样本中缺陷域特征的识别能力.模型的训练成功率有了显著上升，对于纹理样本来说，最好情况下提升了41.3%，物件样本提升了15.1%.纹理样本的显著提升是因为原有的模型训练效果极差（训练成功率20.4%），缺陷区域边缘过于模糊，对比度低.而注意力机制对缺陷区域进行了自学习权重特征提取，增强了模型的特征识别能力.

此外，CBAM模块的放置位置对模型训练效果也有较大影响.从表3中可以看出，自注意力模块越靠近判别器的头部，模型的训练效果会更好.究其原因在于，输入的单张样本只有尺度和位置转换的差异，自注意力模块在特征图的头部，获取的缺陷频域和空间信息特征强度更高，得到的自注意力表达能力更好.

3.4 评估指标

在图像生成任务上来说，常见的评估指标从两个方面着手.一个是生成图像同原始图像的相似度，该数值越小更能表明生成的图像“真实性”越高.另一个是图像中物体的类别丰富度，该指标越大说明模型的效果越好，在同一相似度下，分布越高说明生成模型的表示能力更多样.

在本次实验中使用了以下两种评估指标.

Fréchet Inception Distance（FID［23］），FID是一种用于评估生成模型创建的图像质量指标，较低的分数与较高质量的图像有很高的相关性.

FID度量是两个多维高斯分布之间的平方Wasserstein度量公式（10）为：

因为图像为单张采样生成，所以使用SIFID（Single FID）作为改进后的评估标准.此处计算出的值为训练成功后模型生成的多张图像平均而来.公式（11）为

其中，原真实多张图像和生成图像的数据的统计分布被单样本与其训练后的模型产生的多张生成图像进行一对一分布比较后求和取平均，得到了评估指标SIFID.

NIQE指标［24］是一个客观的评价指标，一般来是提取自然景观中的特征来对测试图像进行测试，这些特征是拟合成一个多元的高斯模型.

在本次实验中，自然景观图像分布被替换为该类别某特定种类特定缺陷图片总的分布，而不再是具体的某单张图像，公式（12）为：

其中v1、v2、Δ1与Δ2分别是样本MVG模型和图像的MVG模型的平均向量和协方差矩阵，详细关于MVG模型介绍见文献［24］.该指标反映生成图像多样性，因此越大越优.

3.5 实验结果分析

实验参数和耗时开销见表4.从表4中可以看出虽然自注意力机制使模型的大小和训练用时同ConSinGAN相比都有增加，但是模型的稳定性和生成效果都有显著提高.模型开销仍显著小于提出单张样本生成的网络SinGAN.

表4 不同模型训练情况表Tab.4 Models Training Table

对于物体类缺陷样本，加入自注意力模块提高模型的生成能力，FID指标有一定的下降，这是因为SSIM模块对相似性的制约.对纹理表面图像来说，SSIM模块显著提高了模型训练的成功率，相比与原始SinGAN模型提升了59%，并且因为该模块对纹理的自适应能力，模型训练出的图像纹理更加丰富.因为对抗生成网络产生的图像是现实生活中不存在的图像，并且该实验为单张样本采样，因此产生的图像有一定的失真.

在纹理样本生成中，单张样本生成的第一个神经网络SinGAN训练成功率很低，并且即使训练成功，模型生成的图像明显失真.木板上的纹理特征学习能力很差，容易产生模糊或者褪色情况，而对较小目标的缺陷部分很难学习到.而ConSinGAN网络具有一定的纹理提取能力但是仍同真实纹理有一定的差异，并且对缺乏对缺陷的学习能力.这主要是因为，这两个网络的设计没有考虑到前景区域对比度低且小的情况，而又缺乏一定前景定位指导能力.

表5为生成模型生成图像较为稳定的4类样本的量化统计情况，每种都生成50张样本进行量化分析.从表5可以看出，从综合四类样本计算得到的参数来看，最终改进的模型相比于SinGAN量化指标SIFID降低了，NIQE提高了.相比于SinGAN模型SIFID降低82.3%，NIQE提高10%.而具体的纹理表现需要在消融对比试验上看出.

表5 不同模型量化指标表Tab.5 Table of quantitative indicators of models

3.6 消融对比实验

在原ConSinGAN实验中，原作者只对风景图像进行了生成，而风景图像具有前景和背景差异较大的特点，因此原模型［8］在公开缺陷样本数据集［22］型训练效果较差.

考虑到缺陷部分提取困难、且要达到端到端的无指导单张样本训练神经网络，因此加入了自注意力网络模块CBAM.

图5中反映了不同的模型训练特定样本产生的图像的典型情况.本文引入的CBAM模块较为成功给缺陷区域加上了权重，使得缺陷不仅没有被池化过程过滤掉，而且成功学习到缺陷的特征，因此能较高概率的产生缺陷，但是仍然存在一定的失真和模糊的情况.此外，图像的边缘因为卷积核边缘填充的原因容易出现明显的四个隐框线的情况.

从图5红色矩形框中可以明显看出，加入了SSIM区域自适应损失函数后，网络的纹理自然程度较高，并且缺陷的锐度明显提高，边缘更加清晰，并有概率出现单缺陷图像训练训练的网络模型，模型可以产生多个缺陷的情况.

本文选取了4个代表性生成图像效果构成图5，并进行如下分析：

图5 生成模型产生的图像比较Fig.5 Examples of images generated by the generation model

（1）第1个样本为带有划痕的木板图片.SinGAN在网络成功完训练的情况下，仍无法学习到表面纹理和缺陷特征，而ConSinGAN网络学习到了大部分纹理和缺陷特征，但是因为训练中的卷积核填充机制，导致图像四周存在不符合原图像的像素虚线.而最终改进的模型不论是纹理真实感还是缺陷的分布都更加自然.

（2）第2个样本为带有水渍的木板图片.SinGAN模型虽然能成功训练，但是生成图像的纹理和缺陷模糊程度很高，实际效果很差.而ConSinGAN在没有加入注意力机制前对缺陷学习能力差.没有加入结构相似度之前，纹理特点不够丰富，而最后采用的改进模型在两者上都有好的表现.

（3）第3个样本为损毁三角管图片.在ConSinGAN网络会出现明显的失真，图像会有较大面积的模糊扭曲.比如三角管主体和背景中黑色孔同色，但是SinGAN难以区分前景和背景，极易导致学习后的网络出现问题.而ConSinGAN因为随机噪声和正则化过程没有受到约束，模型的生成图像有一定的改善但是仍存在部分区域的失真.

（4）第4个样本为内部开裂电缆图片.SinGAN网络学习到的物体结构破坏严重，不论是外壳还是内胶皮都存在不真实的表现（开裂过深，外壳边缘模糊）.ConSinGAN网络无法解决圆形物体的生成约束，而最后改进的模型最终在开裂形变的生成和物体本身特点上取得较为优异的生成能力.

总之，从以上代表性样本示例可以看出，本文改进的两个方式从类间（三极管和背板，电缆外胶、内壳与芯体，缺陷和非缺陷部分）和物体纹理特征（SSIM）两个角度，显著提高了生成图像的真实性.

4 结论

本文将单样本生成引入工业生产缺陷图像样本，解决了工业采样中缺陷频率低导致的图像正负样本量不平衡问题.通过结合结构相似度损失函数和自注意力机制模块，极大改进了原网络无法解决了工业缺陷图像失真和纹理训练困难的问题.

该模型具有训练样本小（单张），训练时间优异（30 min左右）和稳定性较强的优点.实验结果表明本文的改进模型具有良好的性能和效果，具有很好的实用意义.

本文工作还有待于加强和改进，例如对于图5中第四行最右边红色矩形框对应的线缆内层绝缘胶皮，因为其原样图片拍照效果不佳，边缘模糊程度接近与某些随机分布，导致生成器学习效果欠缺，需要在后续工作中考虑到进一步消除缺陷边缘的失真效果.