基于改进生成对抗网络的图像数据增强方法
Image data enhancement method based on improved generative adversarial network
收稿日期: 2022-12-13
基金资助: |
|
Received: 2022-12-13
Fund supported: | 浙江省科技计划资助项目[重点研发(尖兵)项目](2023C01063) |
作者简介 About authors
詹燕(1976—),女,副教授,从事图像处理及智能制造研究.orcid.org/0000-0002-6861-8005.E-mail:
为了提高机器学习模型的精确度,提出基于数据分布拟合、生成式对抗神经网络和图像超分辨率重建的图像数据增强方法. 该方法将最大似然估计和采样算法生成的符合原始数据分布的二维噪声用于对抗训练,克服了在生成模型中传统图像噪声输入随意的问题;采用逐层训练方式生成高分辨率图像,改进高分辨率图像映射困难、参数冗余的缺点. 以轴承滚子表面灰度图像数据增强为例,验证所提方法的有效性. 研究结果表明,所提方法生成的图像质量更优,相比传统方法生成的图像峰值信噪比提高13.07%,结构相似性提高32.40%,弗雷歇初始距离降低37.58%,且数据增强后的模型平均精确度提升7.89%.
关键词:
An image data enhancement method based on data distribution fitting, generative adversarial neural network and image super-resolution reconstruction was proposed to improve the accuracy of machine learning model. The maximum likelihood estimation and sampling algorithm were used to generate two-dimensional noise conforming to the original data distribution for counter-training. The problem of random noise input in traditional image generation models was overcome. Layer by layer training method was used to generate high-resolution images to correct the shortcomings of difficult mapping to high-resolution images. The effectiveness of the proposed method was verified by taking the gray image data enhancement of bearing roller surface. The results showed that the image quality generated by the proposed method was superior. Compared with the image generated by the traditional method, the image peak signal-to-noise ratio was increased by 13.07%, the structural similarity was increased by 32.40%, Fréchet inception distance was reduced by 37.58%, and the average accuracy of the model after data enhancement was increased by 7.89%.
Keywords:
本文引用格式
詹燕, 胡蝶, 汤洪涛, 鲁建厦, 谭健, 刘长睿.
ZHAN Yan, HU Die, TANG Hong-tao, LU Jian-sha, TAN Jian, LIU Chang-rui.
现代零件加工质量检测和设备故障诊断方法逐渐智能化,在各类机器学习和数据挖掘算法中,原始工业数据显得极为重要. 性能良好的算法需要海量且优质的数据支撑,仅用少量数据驱动的模型很难具备较好的泛化能力. 不过,多数情况下难以获取大量的工业数据集,例如变工况条件下标定完整的轴承振动数据集[1].
在图像处理领域,针对工业数据不足的问题,现有的数据增强方法包括对原始图像进行几何变换、随机调整亮度和对比度、添加各类噪声等操作,然而这些方法不能使训练样本的多样性产生质变[2]. 迁移学习[3]可以大幅减少人工标定成本,但是迁移学习模型难以改变其网络结构,且灵活性较差,当训练集和测试集数据分布差异过大时,模型易发生崩溃. 随着生成模型的发展,逐渐出现基于深度玻尔兹曼机(deep Boltzmann machine, DBM)[4]、生成随机网络(generative stochastic network, GSN)[5]、变分自编码器(variational auto-encoding, VAE)[6]、像素递归神经网络(pixel recurrent neural networks, PixelRNN)、像素卷积神经网络(pixel convolution neural networks, PixelCNN)[7-8]和生产式对抗网络(generative adversarial networks, GAN)等数据增强方法. 与VAE相比,GAN不存在偏置,能够更好地拟合真实样本分布;与DBM、GSN和PixelRNN/CNN相比,GAN可一次性生成样本,而不用反复计算马尔可夫链或通过逐个生成像素的方式生成样本. GAN以其良好的图像质量和快速的运行速度,成为了当前数据增强方法的主流研究方向之一.
基于生成器与判别器零和博弈的思想,Goodfellow等[9]提出生成式对抗神经网络,并将其应用于手写数字和人脸图像生成. Mirza等[10]将标签信息输入到GAN的生成器和判别器中,提出条件生成对抗网络(conditional generative adversarial nets, CGAN),通过引入条件信息实现GAN的训练过程可控. Radford等[11]将深度神经网络(deep neural networks,DNN)引入GAN,并用全局池化层替换全连接层,提出深度卷积生成对抗网络(deep convolutional generative adversarial networks,DCGAN). 为了稳定GAN的训练,Arjovsky等[12]提出沃瑟斯坦生成对抗网络(Wasserstein GAN,WGAN). Gulrajani等[13]提出带梯度惩罚项的沃瑟斯坦生成对抗网络(Wasserstein GAN with gradient penalty,WGAN-GP). 传统GAN方法生成的图像数据较为模糊,为了生成高分辨率的图像数据,Berthelot等[14]基于自编码器提出边界均衡生成式对抗网络(boundary equilibrium generative adversarial networks, BEGAN). Karras等[15]提出PG-GAN,通过渐进训练的方式增大生成图像的空间分辨率. Karras等[16-18]借鉴风格迁移网络并通过修改渐进层输入,提出StyleGAN系列对抗网络,实现对隐空间的解耦. 图像超分辨率重建可以利用低分辨率图像经过训练得到高分辨率图像,Dong等[19]将深层卷积网络应用于图像超分辨率重建提出超分辨率卷积神经网络(super-resolution convolutional neural network, SRCNN). 相比通过GAN方法得到高分辨率图像,图像超分辨率重建技术模型结构更简单,且训练过程更稳定.
综上所述,现有图像数据增强方法存在以下问题:1)基于传统GAN方法直接建立向高分辨率图像的映射网络,网络难以工作. 2)图像超分辨率重建技术是在已知低分辨率图像的基础上生成高分辨率图像,模型无法自主生成低分辨率图像. 3)目前基于GAN的模型对于噪声的输入都极为随意,在多数情况下都是直接输入一维随机噪声,导致输入的噪声和样本原始分布差异较大. 网络可训练参数量过多,会影响模型收敛速度. 4)现有的渐进式训练方法能够生成高分辨率图像,但是也带来了训练速度缓慢的问题. 针对上述问题,本研究提出基于分布拟合对抗神经网络的图像数据增强方法. 通过最大似然估计拟合原始样本数据空间分布;根据Box-Muller和马尔科夫链蒙特卡洛(Markov chain Monte Carlo, MCMC)采样算法[20]生成符合原始样本空间分布的随机噪声;结合带条件信息的WGAN-GP和SRCNN提出新的图像数据增强方法,并利用轴承滚子表面缺陷检测数据验证所提方法的可行性和优越性.
1. 基于传统对抗式神经网络的图像数据增强方法
1.1. GAN和CGAN
GAN由生成模型(G)和判别模型(D)2个部分构成,生成模型输入随机噪声
式中:
CGAN的结构以及原理和GAN类似,CGAN在生成模型输入中增加了期望生成数据标签。判别模型输入中增加了真实数据标签,用标签信息控制CGAN的训练过程,使得CGAN能够生成和标签信息对应的数据。CGAN的目标函数为
式中:y为条件信息。
1.2. 基于GAN和CGAN的图像数据增强方法
基于传统对抗式生成网络的图像数据增强方法首先将一维随机噪声和真实图像样本输入到原始GAN或CGAN网络中,同时对生成模型和判别模型进行对抗训练,在模型训练完成后,单独取出生成模型,将一维随机噪声输入生成模型即可生成新的样本,基于传统GAN方法的训练过程如图1所示.
图 1
图 1 基于传统GAN的图像数据增强方法
Fig.1 Image data enhancement method based on traditional GAN
上述2个模型能够生成新的样本,但是都难以训练. 原始GAN和CGAN在训练过程中容易产生模式崩溃和梯度消失问题,并且生成器和判别器的损失值无法反馈模型收敛信息[22]. 一维随机噪声长度较长,会导致全连接层参数过多,造成模型参数冗余,严重影响训练效率.
2. 基于生成对抗神经网络超分辨率重建的图像数据增强方法
为了解决传统GAN用于图像数据增强时存在的模型难以训练、参数大量冗余且噪声输入随意等问题,通过极大似然估计、Box-Muller和MCMC算法生成符合原始数据分布的随机噪声,在此基础上融合WGAN-GP和SRCNN网络,提出基于分布拟合对抗神经网络的图像数据增强方法. 本研究的数据增强方法主要由数据预处理、原始数据分布拟合、随机噪声生成、模型训练及优化和生成样本质量评价等5个步骤组成,框架如图2所示.
图 2
图 2 基于生成对抗神经网络超分辨率重建的图像数据增强方法流程图
Fig.2 Flowchart of image data enhancement method based on generative adversarial neural network super-resolution reconstruction
2.1. WGAN和WGAN-GP
生成器随机初始化后的生成分布很难与真实分布有不可忽略的重叠,因此WGAN用Wasserstein距离作为等价优化的距离衡量,进而同时解决稳定训练和进程指标的问题[12]. Wasserstein距离计算公式为
式中:
Wasserstein距离计算公式的对偶形式为
式中:
WGAN的目标函数为
式中:
式中:
2.2. 数据预处理
对工业相机采集的少量原始图像数据进行人工标定,获得图像对应标签;用单层卷积层对原始图像数据集进行降采样,将原始图像尺寸缩小到所需尺寸. 目的是减少原始图像数据的空间冗余,提高模型的计算效率. 降采样输出图像尺寸计算公式为
式中:
2.3. 原始数据分布拟合
通过最大似然估计法对降采样之后的数据进行分布拟合,目的是获得原始数据分布的概率密度函数. 首先生成图像灰度直方图,根据灰度直方图初步判断数据近似服从的分布;然后分别将近似服从分布和常用分布的概率密度函数构建成最大似然函数;最后求解得到概率密度函数待估计参数值. 最大似然估计法计算方法为
式中:
2.4. 随机噪声生成
通过Box-Muller和MCMC采样算法,根据原始数据分布拟合得到的概率密度函数生成符合原始数据分布的随机噪声,具体方法步骤如图3所示.
图 3
2.4.1. Box-Muller变换
在很多情况下累积分布函数的反函数无法直接获得. Box-Muller变换通过对目标分布的联合概率密度函数进行三角换元,分别求出换元后联合概率密度函数关于
式中:
2.4.2. Metropolis-Hastings抽样算法
图 4
图 4 Metropolis-Hastings抽样算法流程图
Fig.4 Flowchart of Metropolis-Hastings sampling algorithm
假设
2.5. 建立生成式对抗网络和SRCNN网络
图 5
图 6
模型生成器损失函数为
SRCNN损失函数选用均方误差(mean-square error,MSE)损失函数,计算式为
式中:
2.6. 生成图像质量评价
图像峰值信噪比(peak signal-to-noise ratio, PSNR)和结构相似性(structural similarity, SSIM)能对生成图像质量进行量化. PSNR是反映图像对应像素点间误差的客观评价方法,SSIM则是从亮度、对比度和结构3个方面衡量生成图像质量,更符合人眼视觉感受. 具体计算方法为
式中:
弗雷歇初始距离(Fréchet inception distance,FID)作为衡量真实图像和生成图像之间特征距离的指标之一,既能够表示图像的多样性,也能被用来评价生成图像的质量.
式中:
3. 实际案例分析
由于轴承滚子各个表面图像像素分布差异较大,且轴承表面样本采集存在一定难度,须先人工擦拭滚子各表面,再通过采集设备逐一采集,人力成本耗费较大,比较贴合实际工业生产中样本采集现状. 以轴承滚子表面灰度图像数据增强为例,验证所提方法的有效性,原始样本灰度图像均通过CCD工业相机采集获得,样本图像分为3类,分别为轴承滚子侧面、倒角、端面灰度图像,数据集中3类表面图像数量分别为351、468、471,总计1290. 输入判别器的图像需要降采样,输入生成器的是二维噪声. 首先将维度为(1 920, 1 200, 1)的原始图像裁剪为(1 920, 1 184, 1),然后用卷积核数量为1的单层卷积层对裁剪后图像数据进行降采样,分别降采样为(480, 296, 1)和(240, 148, 1),并将滚子侧面、倒角、端面图像标签记为0、1、2. 维度为(240, 148, 1)的图像用于训练WGAN-GP网络,维度为(480, 296, 1)的图像作为训练SRCNN模型的高分辨率图像,训练好的生成对抗网络生成的800张维度为(240, 148, 1)的图像作为训练SRCNN模型的低分辨率图像,其中侧面、端面、倒角面图像分别为200、300、300张.
标签相同的图像灰度值分布接近一致,随机从3类图像中各抽取一张图像,生成3类轴承滚子图像的灰度直方图,初步判断图像近似符合的分布类型,降采样后图像数据以及对应的灰度直方图如图7所示. 图中,
图 7
以滚子端面为例,拟合图像原始数据分布,首先找出图像灰度值可能服从的分布,写出其待估计参数的概率密度函数,根据式(9)~(11)构建并求解最大似然函数,得到待估计参数值,通过计算真实值与拟合值之间的误差平方和(sum of squares for error, SSE)评价拟合程度. SSE计算方法为
式中:
表 1 轴承滚子端面灰度直方图拟合结果
Tab.1
拟合分布名称 | 估计参数 | SSE | ||
α | β | ε | ||
卡方分布 | 2.848 | 45.687 | 2.848 | 0.004 551 |
伽马分布 | −20.172 | 8.656 | 9.378 | 0.004 411 |
高斯分布 | 65.852 | 31.368 | 0 | 0.004 014 |
T分布 | 66.366 | 17.124 | 2.734 | 0.002 338 |
拉普拉斯分布 | 68.000 | 19.693 | 0 | 0.002 049 |
韦布尔分布 | 69.000 | 20.150 | 0.926 | 0.002 051 |
柯西分布 | 69.070 | 11.861 | 0 | 0.002 038 |
图 8
图 8 轴承滚子端面灰度直方图拟合效果
Fig.8 Gray histogram fitting effect of bearing roller end face
图 9
图 9 轴承滚子各表面灰度直方图的最优分布拟合效果
Fig.9 Optimal distribution fitting effect of gray histogram of bearing roller surface
表 2 轴承滚子倒角面及侧面灰度直方图的拟合结果
Tab.2
图像名称 | 估计参数 | SSE | ||
α | β | ε | ||
倒角面 | 8.305 | 0.475 | 2.000 | 0.046284 |
侧面(峰1) | 8.360 | 4.882 | 0 | 0.007293 |
侧面(峰2) | 141.441 | 21.260 | 0 | 0.004029 |
侧面(峰3) | 254.950 | 1.754 | 0 | 0.023585 |
式中:α0、α1、α2为位置参数,β0、β1、β2为尺度参数.
倒角面图像最优拟合分布为三参数韦布尔分布,概率密度函数为
已知3类样本图像分布的概率密度函数,根据Box-Muller变换和MH采样算法生成符合样本分布的随机噪声,端面数据为柯西分布,直接通过Box-Muller变换求得概率密度函数关于R和
式中:
图 10
图 10 Box-Muller变换生成的端面一维随机噪声分布
Fig.10 One-dimensional random noise distribution of end face generated by Box-Muller transform
图 11
图 12
图 12 采样数据分布与原始数据分布的对比图
Fig.12 Comparison between distribution of sampled data and distribution of raw data
将使用的采样方法构建成噪声生成模块,外接在生成式对抗网络前端,为网络生成符合原始样本分布的噪声. 当搭建WGAN-GP网络时,判别器最后一层为未激活的全连接层,设置梯度惩罚权重参数
图 13
图 14
由图13可知,传统CGAN网络难以训练,在训练过程中会出现严重的梯度消失问题,导致模型经过22000个批次训练后依旧无法有效收敛. WGAN-GP相比CGAN收敛效果较好,经过7000个批次训练后,模型已经具备拟合原始图像基本轮廓的能力,但是直接向高分辨率图像建立映射关系,生成图像质量较差. 所提方法采用逐层训练的方式,避免了生成对抗网络直接向高分辨率图像映射的弊端,而是通过SRCNN提升图像分辨率,相比单独使用CGAN和WGAN-GP模型,本研究所提模型能够获得最优的生成图像质量.
由图14可知,CGAN网络由于判别器过于强大,经过600次迭代后,判别器接近收敛,导致模型优化饱和,生成器无法学到有用的信息,输出图片质量较差.
图 15
图 16
各模型生成图像质量评价指标如表3所示. CGAN模型由于生成器无法学习到有效信息,生成图像质量较差,本研究所提模型相比WGAN-GP模型,生成图像的信噪比提高13.07%,平均结构相似性提高32.40%,平均FID降低37.58%.
表 3 各模型生成图像质量评价表
Tab.3
模型 | PSNR/dB | SSIM | FID | ||||||||
侧面 | 倒角面 | 端面 | 侧面 | 倒角面 | 端面 | 侧面 | 倒角面 | 端面 | |||
CGAN | 48.043 3 | 54.327 9 | 52.459 2 | 0.184 2 | 0.288 9 | 0.248 0 | 472.899 4 | 343.297 9 | 372.268 5 | ||
WGAN-GP | 57.523 0 | 62.487 9 | 53.688 9 | 0.634 2 | 0.831 7 | 0.396 4 | 363.993 2 | 166.245 8 | 207.444 2 | ||
WGAN-GP+SR | 63.543 6 | 68.508 5 | 64.347 9 | 0.807 0 | 0.941 1 | 0.717 6 | 224.5837 | 132.936 6 | 102.986 2 |
图 17
图 17 不同标准差噪声及对应生成图像
Fig.17 Different standard deviation noise and corresponding image generation
图18为不同采样方式在WGAN-GP网络训练过程中的FID变化曲线,图中,MAP为平均精确度. 所提的预先对样本图像分布进行拟合的随机采样方法,相比直接采用均匀分布和正态分布生成噪声的方法更利于模型训练.
表 4 各模型生成图像质量评价表
Tab.4
噪声标准差 | PSNR/dB | SSIM | FID |
0.168 6 | 56.580 9 | 0.599 2 | 312.831 4 |
0.215 5 | 56.375 6 | 0.371 3 | 406.542 9 |
0.784 0 | 55.263 4 | 0.165 7 | 430.592 1 |
图 18
图 18 轴承滚子数据集不同采样方式训练过程FID曲线图
Fig.18 Training process FID curve of bearing roller data set with different sampling methods
如图19所示为WGAN-GP+SR和WGAN-GP图像的对比. 可以看出,相比直接用对抗网络生成的图像,超分辨率重建后的图像质量明显提升.
图 19
图 19 WGAN-GP+SR与WGAN-GP图像对比
Fig.19 Comparison of WGAN-GP+SR and WGAN-GP images
如图20所示为对抗网络数据增强前后Yolov5目标检测模型训练过程平均精确度(MAP)曲线对比图,MAP越高,检测效果越好. 从图中可知,数据增强前和数据增强后MAP峰值分别为0.532和0.574,对抗网络数据增强后模型MAP提高7.89%.
图 20
图 20 Yolov5目标检测模型训练过程MAP曲线图
Fig.20 Map diagram of Yolov5 target detection model training process
如图21所示,对比3组数据增强前后对轴承滚子缺陷检测的结果. 可以看出,第1组中Hump类缺陷在数据增强前被误检为ResidualPit类,而第2组和第3组中的Hump类缺陷在数据增强前则皆被漏检,证明通过对抗网络数据增强能够降低检测模型的误检率和漏检率.
图 21
4. 结 语
对图像数据增强方法做了深入研究,提出基于生成对抗神经网络超分辨率重建的图像数据增强方法,通过案例分析和模型对比验证了所提方法具有最优的生成图像质量. 通过拟合原始数据分布构建噪声生成模块,为模型输入符合数据分布的二维噪声,有效降低了模型的参数冗余度,提高了模型生成图像质量的计算经济性. 融合WGAN-GP和SRCNN模型提出一种复合模型,避免生成对抗网络直接向高分辨率图像建立映射关系的弊端. 本研究所提出的生成对抗网络结合图像超分重建思想所用模型并不固定,下一步可以尝试不同的模型组合,例如BEGAN+SRCNN.
参考文献
基于特征迁移学习的变工况下滚动轴承故障诊断方法
[J].DOI:10.13334/J.0258-8013.PCSEE.180130 [本文引用: 1]
Fault diagnosis method of a rolling bearing under variable working conditions based on feature transfer learning
[J].DOI:10.13334/J.0258-8013.PCSEE.180130 [本文引用: 1]
基于二维灰度图的数据增强方法在电机轴承故障诊断的应用研究
[J].DOI:10.13334/j.0258-8013.pcsee.200834 [本文引用: 1]
Research on the application of the data augmentation method based on 2D gray pixel images in the fault diagnosis of motor bearing
[J].DOI:10.13334/j.0258-8013.pcsee.200834 [本文引用: 1]
A survey on transfer learning
[J].
Generative adversarial networks
[J].
Alias-free generative adversarial networks
[J].
Image super-resolution using deep convolutional networks
[J].
Twister generator of random normal numbers by box-muller model
[J].
基于隐变量后验生成对抗网络的不平衡学习
[J].DOI:10.16183/j.cnki.jsjtu.2019.264 [本文引用: 1]
Unbalanced learning of generative adversarial network based on latent posterior
[J].DOI:10.16183/j.cnki.jsjtu.2019.264 [本文引用: 1]
Sampling methods for solving bayesian model updating problems: a tutorial
[J].DOI:10.1016/j.ymssp.2021.107760 [本文引用: 1]
/
〈 |
|
〉 |
