浙江大学学报(工学版), 2023, 57(10): 1998-2010 doi: 10.3785/j.issn.1008-973X.2023.10.009

计算机技术、自动化技术

基于改进生成对抗网络的图像数据增强方法

詹燕,, 胡蝶, 汤洪涛, 鲁建厦, 谭健, 刘长睿

浙江工业大学 机械工程学院,浙江 杭州 310023

Image data enhancement method based on improved generative adversarial network

ZHAN Yan,, HU Die, TANG Hong-tao, LU Jian-sha, TAN Jian, LIU Chang-rui

College of Mechanical Engineering, Zhejiang University of Technology, Hangzhou 310023, China

收稿日期: 2022-12-13  

基金资助: 浙江省科技计划资助项目[重点研发(尖兵)项目](2023C01063)

Received: 2022-12-13  

Fund supported: 浙江省科技计划资助项目[重点研发(尖兵)项目](2023C01063)

作者简介 About authors

詹燕(1976—),女,副教授,从事图像处理及智能制造研究.orcid.org/0000-0002-6861-8005.E-mail:yzhan@zjut.edu.cn , E-mail:yzhan@zjut.edu.cn

摘要

为了提高机器学习模型的精确度,提出基于数据分布拟合、生成式对抗神经网络和图像超分辨率重建的图像数据增强方法. 该方法将最大似然估计和采样算法生成的符合原始数据分布的二维噪声用于对抗训练,克服了在生成模型中传统图像噪声输入随意的问题;采用逐层训练方式生成高分辨率图像,改进高分辨率图像映射困难、参数冗余的缺点. 以轴承滚子表面灰度图像数据增强为例,验证所提方法的有效性. 研究结果表明,所提方法生成的图像质量更优,相比传统方法生成的图像峰值信噪比提高13.07%,结构相似性提高32.40%,弗雷歇初始距离降低37.58%,且数据增强后的模型平均精确度提升7.89%.

关键词: 图像数据增强 ; 分布拟合 ; 采样算法 ; 生成式对抗网络 ; 图像超分辨率重建

Abstract

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: image data enhancement ; distribution fitting ; sampling algorithm ; generative adversarial neural network ; image super-resolution reconstruction

PDF (5142KB) 元数据 多维度评价 相关文章 导出 EndNote| Ris| Bibtex  收藏本文

本文引用格式

詹燕, 胡蝶, 汤洪涛, 鲁建厦, 谭健, 刘长睿. 基于改进生成对抗网络的图像数据增强方法. 浙江大学学报(工学版)[J], 2023, 57(10): 1998-2010 doi:10.3785/j.issn.1008-973X.2023.10.009

ZHAN Yan, HU Die, TANG Hong-tao, LU Jian-sha, TAN Jian, LIU Chang-rui. Image data enhancement method based on improved generative adversarial network. Journal of Zhejiang University(Engineering Science)[J], 2023, 57(10): 1998-2010 doi:10.3785/j.issn.1008-973X.2023.10.009

现代零件加工质量检测和设备故障诊断方法逐渐智能化,在各类机器学习和数据挖掘算法中,原始工业数据显得极为重要. 性能良好的算法需要海量且优质的数据支撑,仅用少量数据驱动的模型很难具备较好的泛化能力. 不过,多数情况下难以获取大量的工业数据集,例如变工况条件下标定完整的轴承振动数据集[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个部分构成,生成模型输入随机噪声 $z $,用于拟合真实样本数据分布,判别模型接收生成样本 $G(z) $以及真实样本 $x $,用于判别生成样本 $G(z) $是否来自于真实样本x。生成模型的目的是最大化 $D(G(z)) $,即最大化判别模型将 $G(z) $判别为真实样本的可能性,而判别模型的目的是最小化 $D(G(z)) $,即最小化将 $G(z) $判别为真实样本的可能性。在经过多次对抗训练后,生成器和判别器达到纳什平衡[21]。GAN的目标函数为

$ \begin{split} & \mathop {\min }\limits_G \;\mathop {\max }\limits_D \;V(D,G) = {E_{x \sim {P_{{\rm{r}}}}(x)}}[\ln \; D(x)]+ \\ &\quad {E_{z \sim {P_z}(z)}}[\ln \; (1 - D(G(z)))]. \\ \end{split} $

式中: $V(D,G) $为优化目标函数,由生成器和判别器2个部分组成; $E $为分布函数的数学期望; $x $为真实数据; ${P_{{\rm{r}}}}(x)$为原始数据分布; $z $为随机噪声; ${P_z}(z) $为随机噪声数据分布; $D(x) $为判别模型判别结果; $G(z) $为生成模型输出样本。

CGAN的结构以及原理和GAN类似,CGAN在生成模型输入中增加了期望生成数据标签。判别模型输入中增加了真实数据标签,用标签信息控制CGAN的训练过程,使得CGAN能够生成和标签信息对应的数据。CGAN的目标函数为

$\begin{split} & \mathop {\min }\limits_G \; \mathop {\max }\limits_D \; V(D,G) = {E_{x \sim {P_{{\rm{r}}}}(x)}}[\ln \; D(x\left| y \right.)] +\\ &\quad {E_{z \sim {P_z}(z)}}[\ln \;(1 - D(G(z\left| y \right.)))]. \end{split} $

式中: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距离计算公式为

$ W({P_{\rm{r}}},{P_{\rm{g}}}) = \mathop {\inf }\limits_{\gamma \sim \prod ({P_{\rm{r}}},{P_{\rm{g}}})} {E_{(x,y) \sim \gamma }}[||x - y||]. $

式中: $ {P_{\rm{r}}} $为真实数据分布; $ {P_{\rm{g}}} $为生成数据分布; $\prod ({P_{\rm{r}}},{P_{\rm{g}}})$$ {P_{\rm{r}}} $$ {P_{\rm{g}}} $组合得到的所有可能的联合分布集合;对于每个可能的联合分布 $\gamma $可以从中采样 $(x,y) \sim \gamma $得到样本 $x$$y$$\left\| {x - y} \right\|$为这对样本的距离.

Wasserstein距离计算公式的对偶形式为

$ \begin{gathered} W({P_{\rm{r}}},{P_{\rm{g}}}) = \frac{1}{K}\mathop {\sup }\limits_{{{\left\| f \right\|}_{L \leqslant K}}} \;\Biggr\{E{}_{x \sim {P_{\rm{r}}}}[f(x)] - {E_{x \sim {P_{\rm{g}}}}}[f(x)]\Biggr\}. \end{gathered} $

式中: $\sup $为上确界,即最小上界; $f$为一个连续函数; ${\left\| f \right\|_{L \leqslant K}}$表示 $f$必须满足利普希茨连续条件(Lipschitz continuity),即存在一个常数 $K \geqslant 0$使得在定义域内 $\left| {f({x_1}) - f({x_2})} \right| \leqslant K\left| {{x_1} - {x_2}} \right|$.

WGAN的目标函数为

$ V(G,D) = \mathop {\max }\limits_{D \in 1 - Lipschitz} \{ {E_{x \sim {P_{\rm{r}}}}}[D(x)] - {E_{x \sim {P_{\rm{g}}}}}[D(x)]\} . $

式中: $D \in 1 - {\rm{Lipschitz}}$为判别器 $D$必须满足利普希茨连续条件,即满足 $\left| {D({x_1}) - D({x_2})} \right| \leqslant \left| {{x_1} - {x_2}} \right|$. 为了满足这一条件,WGAN在每步迭代完成后将判别器 $D$的参数截断在 $\left[ { - c,c} \right]$,但是这种优化策略容易导致参数取值极端化、梯度消失和梯度爆炸. 针对这个问题,WGAN-GP通过对目标函数加入梯度惩罚项间接实现利普希茨连续条件,WGAN-GP的目标函数为

$ \begin{split} V(G,D) =& \max \;\{ {E_{x \sim {P_{\rm{r}}}}}[D(x)] - {E_{x \sim {P_{\rm{g}}}}}[D(x)] -\\ & \lambda {E_{\hat x \sim {P_{{{\rm{PEN}}}}}}}[{({\left\| {{\nabla _{\hat x}}D(\hat x)} \right\|_2} - 1)^2}]\} , \end{split} $

$ \left. \begin{gathered} \hat x = \varepsilon {x_{\rm{r}}}+(1 - \varepsilon ){x_{\rm{g}}}, \\ {x_{\rm{r}}} \in {P_{\rm{r}}},{x_{\rm{g}}} \in {P_{\rm{g}}},\varepsilon \in {\rm{Uniform}}[0,1.0]. \\ \end{gathered} \right\} $

式中: $\lambda $为梯度惩罚权重, $ {x_{\rm{r}}} $$ {x_{\rm{g}}} $分别为真实分布 $ {P_{\rm{r}}} $和生成分布 $ {P_{\rm{g}}} $中的数据, $\hat x$为随机噪声在 $ {x_{\rm{r}}} $$ {x_{\rm{g}}} $的连线上随机插值采样得到数据, ${P_{{{\rm{PEN}}}}}$为采样得到数据的集合, $\varepsilon $为服从0~1.0均匀分布的随机数, $ {\left\| {{\nabla _{\hat x}}D(\hat x)} \right\|_2} $为判别器 $D$梯度的 ${L_2}$范数.

2.2. 数据预处理

对工业相机采集的少量原始图像数据进行人工标定,获得图像对应标签;用单层卷积层对原始图像数据集进行降采样,将原始图像尺寸缩小到所需尺寸. 目的是减少原始图像数据的空间冗余,提高模型的计算效率. 降采样输出图像尺寸计算公式为

$ {O_{{\rm{out}}}} = \left( {{{{I_{{\rm{in}}}} - {k_{\rm{c}}}+2V}}} \right)/{d} +1. $

式中: $ {O_{{\rm{out}}}} $为输出图像尺寸, $ {I_{{\rm{in}}}} $为输入图像尺寸, $ {k_{\rm{c}}} $为卷积核尺寸, $V$为边界填充参数, $d$为滑移步长. 输出图像通道数等于卷积核数量.

2.3. 原始数据分布拟合

通过最大似然估计法对降采样之后的数据进行分布拟合,目的是获得原始数据分布的概率密度函数. 首先生成图像灰度直方图,根据灰度直方图初步判断数据近似服从的分布;然后分别将近似服从分布和常用分布的概率密度函数构建成最大似然函数;最后求解得到概率密度函数待估计参数值. 最大似然估计法计算方法为

$ L(\theta ) = L({x_1},{x_2},\cdots,{x_n}|\theta ) = \mathop \prod \limits_{i = 1}^n f({x_i}|\theta ) \text{,} $

$ \ln \; L(\theta ) = \sum\limits_{i = 1}^n {\ln \; f({x_i}|\theta )} \text{,} $

$ \begin{gathered} \hat \theta = \arg \mathop {\max }\limits_\theta \; \sum\limits_{i = 1}^n {\ln \; L(\theta )} = \arg \mathop {\max }\limits_\theta \; \sum\limits_{i = 1}^n {\ln \; f({x_i}|\theta )} . \end{gathered} $

式中: $ L(\theta ) $为最大似然函数, $ \theta $估计参数, $ f({x_i}|\theta ) $为概率密度函数, $ {x_i} $为样本值, $ \hat \theta $为极大似然函数估计值.

2.4. 随机噪声生成

通过Box-Muller和MCMC采样算法,根据原始数据分布拟合得到的概率密度函数生成符合原始数据分布的随机噪声,具体方法步骤如图3所示.

图 3

图 3   生成随机噪声流程图

Fig.3   Flowchart of generating random noise


2.4.1. Box-Muller变换

在很多情况下累积分布函数的反函数无法直接获得. Box-Muller变换通过对目标分布的联合概率密度函数进行三角换元,分别求出换元后联合概率密度函数关于 $R$$ \theta $的边缘分布函数,再对其求反函数,最后将服从均匀分布的随机变量映射到所求分布上,得到服从目标分布的随机噪声. Box-Muller变换计算式如下:

$ \iint\limits_\varOmega {f(x,y){{\rm{d}}x}{{\rm{d}}y} = }\iint\limits_{{\varOmega ^{'}}} {f(R\cos \; \theta ,R\sin \; \theta )R{{\rm{d}}R}{{\rm{d}}\theta }}. $

$ {F_R}(R \leqslant r) = \int_0^{2\text{π} } {\int_0^r {f(R\cos \; \theta ,R\sin \; \theta )R{{\rm{d}}R}{{\rm{d}}\theta }} } ;r > 0. $

$ {F_\theta }(\theta \leqslant \phi ) = \int_0^\phi {\int_0^{+\infty } {f(R\cos \; \theta ,R\sin \; \theta )R{{\rm{d}}R}{{\rm{d}}\theta }} } ;0 < \phi \leqslant 2\text{π} . $

$ R = F_R^{ - 1}({U_1}),\theta = F_\theta ^{ - 1}({U_2});{U_1},{U_2} \sim U(0,1.0). $

$ x = F_R^{ - 1}({U_1})\cos \;[F_\theta ^{ - 1}({U_2})]. $

$ y = F_R^{ - 1}({U_1})\sin \;[F_\theta ^{ - 1}({U_2})]. $

式中: $ f(x,y) $$x$ $y$的联合概率密度函数,且 $x$ $y$相互独立; $ {F_R} $$ F_R^{ - 1} $为关于 $R$的边缘分布函数和其反函数; $ {F_\theta } $$ F_\theta ^{ - 1} $为关于 $\theta $的边缘分布函数和其反函数; $ {U}_{1}、{U}_{2} $为服从0~1.0均匀分布的随机数; $x$$y$服从均值为0,方差为1.0的高斯分布.

2.4.2. Metropolis-Hastings抽样算法

对于目标分布的累积分布函数或关于 $R$$ \theta $的边缘分布函数反函数不可求的情况,可以通过Metropolis-Hastings(MH)抽样算法得到符合目标分布的随机噪声[23]. MH算法是MCMC采样算法中常用的一种,通过构建马尔可夫链,找到满足细致平稳条件的状态转移矩阵使马尔可夫链趋于平稳分布,算法流程如图4所示.

图 4

图 4   Metropolis-Hastings抽样算法流程图

Fig.4   Flowchart of Metropolis-Hastings sampling algorithm


假设 $P(x)$为目标概率分布,对于 $P(x)$构建一个关于样本 $x$的马氏链,其状态转移矩阵记为 ${\boldsymbol{Q}}$. 在MH算法中 ${\boldsymbol{Q}}$即为给定的建议分布, $x'$为按建议分布随机抽取的候选状态, ${x_{t - 1}} = x'$$\alpha (x,x')$为接收概率,MH抽样算法根据 $\alpha (x,x')$决定是否用 $x'$更新 ${x_t}$$t$为当前时刻, $u$为服从0~1.0均匀分布的随机数,具体计算方法为

$ \alpha (x,x') = \min \; \left\{ {1,\frac{{P(x')Q(x',x)}}{{P(x)Q(x,x')}}} \right\}. $

$ {x_t} = \left\{ {\begin{array}{*{20}{c}} {x',\;u \leqslant \alpha (x,x');} \\ {x,\;u > \alpha (x,x').} \end{array}} \right. $

2.5. 建立生成式对抗网络和SRCNN网络

通过融合带条件信息的WGAN-GP和SRCNN建立生成式对抗网络模型,可以使模型能够按照给定标签信息生成对应的样本图像,且减少模式崩溃、梯度消失问题的产生. 本研究将WGAN-GP输出的低分辨率图像输入到预训练好的SRCNN网络中,经过超像素重建后得到高分辨率图像. 为了探寻输入随机噪声分布对生成式对抗网络模型的影响,构建一个随机噪声生成模块,为模型生成符合样本分布的随机噪声,WGAN-GP网络架构如图5所示,模型整体结构如图6所示. 模型判别器损失函数为

图 5

图 5   WGAN-GP网络架构图

Fig.5   Architecture diagram of WGAN-GP network


图 6

图 6   WGAN-GP+SR模型整体结构图

Fig.6   Overall structure diagram of WGAN-GP+SR model


$ \begin{split} L_D = &\min \; \Big\{ {E_{z \sim {P_z}(z)}}[D(x|y)] \Big.- {E_{x:{P_{\rm{r}}}(x)}}[D(x|y)] + \\ & \left.\lambda {E_{\hat x:{P_{{{\rm{PEN}}}}}}}\left[{\left({\left\| {{{\tilde N}_{\hat x}}D(\hat x)} \right\|_2} - 1\right)^2}\right]\right\} . \end{split} $

模型生成器损失函数为

$ L_G = \min \; \{ - {E_{z \sim {P_z}(z)}}[D(x|y)]\} . $

SRCNN损失函数选用均方误差(mean-square error,MSE)损失函数,计算式为

$ {\rm{MSE}} = \frac{1}{n}\sum\limits_{i = 1}^n {{{\left\| {f({x_i}) - {y_i}} \right\|}^2}} . $

式中: $n$为图像像素的个数, ${y_i}$为真实标签, $ f({x_i}) $为网络输出.

2.6. 生成图像质量评价

图像峰值信噪比(peak signal-to-noise ratio, PSNR)和结构相似性(structural similarity, SSIM)能对生成图像质量进行量化. PSNR是反映图像对应像素点间误差的客观评价方法,SSIM则是从亮度、对比度和结构3个方面衡量生成图像质量,更符合人眼视觉感受. 具体计算方法为

$ {\rm{PSNR}}(X,Y) = 10\times \lg \;\left[ {\frac{{{{({2^b} - 1)}^2}}}{{{\rm{MSE}}(X,Y)}}} \right], $

$ \begin{split} & {\rm{SSIM}}(X,Y) = \\ &\qquad \frac{{(2{\mu _X}{\mu _Y}+{C_1})(2{\sigma _X}{\sigma _Y}+{C_2})({\sigma _{X,Y}}+{C_3})}}{{(\mu _X^2+\mu _Y^2+{C_1})(\sigma _X^2+\sigma _Y^2+{C_2})({\sigma _X}{\sigma _Y}+{C_3})}}. \end{split} $

式中: $ {\rm{MSE}}(X,Y) $为生成图像 $ X $与真实图像 $Y$之间的均方误差; $ \;{\mu }_{X}、{\mu }_{Y}、{\sigma }_{X}、{\sigma }_{Y}、{\sigma }_{X,Y} $分别为 $ X、Y $的局部均值、方差和协方差; $b$为每像素比特数; $ {C}_{1}= {({K}_{1} L)}^{2}、{C}_{2}={({K}_{2} L)}^{2}、{C}_{3}={C}_{3}/2 $,一般情况下 $ {K}_{1}=0.01、{K}_{2}=0.01、L=255.00 $.

弗雷歇初始距离(Fréchet inception distance,FID)作为衡量真实图像和生成图像之间特征距离的指标之一,既能够表示图像的多样性,也能被用来评价生成图像的质量.

$ \begin{gathered} {\rm{FID}}(x,g) = {{T}}_{\text{r}}\left({\bf{Cov}}_x + {\bf{Cov}}_g - 2\sqrt {{\bf{Cov}}_x {{\bf{Cov}} {_g} } } \right) {\text+}{\left\| {{\mu _x} - {\mu _g}} \right\|^2}. \end{gathered} $

式中: $x$为真实图像, $g$为生成图像, $ \;{{\boldsymbol{\mu}} _x} $${\bf{Cov}} {_x} $分别为真实图像的特征向量的平均值和协方差矩阵, $ \;{{\boldsymbol{\mu}} _g} $$ {\bf{Cov}} {_g} $分别为生成图像的特征向量的平均值和协方差矩阵, ${{T}}_{\text{r}}$为矩阵对角线上元素的总和. 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所示. 图中, ${{\rm{GV}}} $为灰度值, $P$为概率. 滚子侧面图像灰度值集中分布在3个区域,为多峰分布;滚子倒角面图像灰度值主要分布在0~50,为偏态分布;滚子端面图像灰度值分布较为均匀,考虑近似服从正态、柯西或拉普拉斯等分布.

图 7

图 7   降采样后图像数据及其灰度直方图

Fig.7   Image data and gray histogram after down-sampling


以滚子端面为例,拟合图像原始数据分布,首先找出图像灰度值可能服从的分布,写出其待估计参数的概率密度函数,根据式(9)~(11)构建并求解最大似然函数,得到待估计参数值,通过计算真实值与拟合值之间的误差平方和(sum of squares for error, SSE)评价拟合程度. SSE计算方法为

$ {\rm{SSE}} = \sum\limits_{i = 1}^n {{{({{\hat y}_i} - {y_i})}^2}} . $

式中: $\hat y_i $$y_i $分别为真实值和拟合值.

滚子端面图像可能服从的分布拟合结果及其SSE如表1所示,表中α为位置参数,β为尺度参数,ε为形状参数. SSE越小说明拟合误差越小,拟合效果越好. 由表1可知,用柯西分布拟合滚子端面灰度直方图拟合误差最小,即柯西分布为滚子端面图像的最优拟合分布,估计位置参数α = 69.070,尺度参数β = 11.861,即端面图像灰度值分布概率密度函数为

表 1   轴承滚子端面灰度直方图拟合结果

Tab.1  Gray histogram fitting results of bearing roller end face

拟合分布名称 估计参数 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

新窗口打开| 下载CSV


$ f(x;\alpha,\beta ) = \frac{1}{\text{π} }\left[\frac{{11.861}}{{{{(x - 69.07)}^2}+{{11.861}^2}}}\right]. $

滚子端面图像灰度直方图各分布拟合效果以及最优拟合分布如图8所示. 同理可得轴承侧面及倒角面最优拟合分布,其拟合效果如图9所示,估计参数和SSE如表2所示,侧面图像最优拟合分布为多峰柯西分布. 由表2可知,侧面图像最优拟合分布概率密度函数为

图 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  Fitting results of gray histogram of bearing roller chamfering and side surface

图像名称 估计参数 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

新窗口打开| 下载CSV


$ \begin{split} f(x;{\alpha_0},{\beta _0},{\alpha_1},{\beta _1},{\alpha_2},{\beta _2}) =& \frac{1}{\text{π} }\left[\frac{{4.882}}{{{{(x - 8.36)}^2}+{{4.882}^2}}}\right]+ \\ & \frac{1}{\text{π} }\left[\frac{{21.26}}{{{{(x - 141.441)}^2}+{{21.26}^2}}}\right] +\\ & \frac{1}{\text{π} }\left[\frac{{1.754}}{{{{(x - 254.95)}^2}+{{1.754}^2}}}\right]. \end{split} $

式中:α0α1α2为位置参数,β0β1β2为尺度参数.

倒角面图像最优拟合分布为三参数韦布尔分布,概率密度函数为

$ f(x;\alpha ,\beta ,\varepsilon ) = \frac{{0.475}}{2}{ \left( \frac{{x - 8.305}}{2} \right) ^{ - 0.525}}{{\rm{exp}}\;{ - {{ \left( \frac{{x - 8.305}}{2} \right) }^{0.475}}}}. $

已知3类样本图像分布的概率密度函数,根据Box-Muller变换和MH采样算法生成符合样本分布的随机噪声,端面数据为柯西分布,直接通过Box-Muller变换求得概率密度函数关于R$\theta $边缘分布函数的反函数. 由式(12)~(17)得到端面图像Box-Muller变换后结果如下:

$ x = \frac{{11.861\sqrt {2{U_1} - U_1^2} }}{{1 - {U_1}}}\cos\; (2\text{π} {U_2})+69.07, $

$ y = \frac{{11.861\sqrt {2{U_1} - U_1^2} }}{{1 - {U_1}}}\sin \;(2\text{π} {U_2})+69.07. $

式中: $ {U}_{1}、{U}_{2} $为服从0~1.0均匀分布的随机数.

根据Box-Muller变换结果得到的一维随机噪声分布如图10所示,图中 $Q$为采样频数. 侧面和倒角面图像概率密度函数形式较为复杂,直接通过MH采样算法生成符合原始数据分布的噪声,设置MH抽样算法 $f(x)$为侧面和倒角面的概率密度函数, $q(x)$均为0~1.0均匀分布,燃烧期M=10 000,迭代步数N=19 000. MH采样算法一维采样结果如图11所示. 3类样本图像二维采样结果与原始分布对比如图12所示. 图中,左侧为采样数据分布,右侧为原始数据分布.

图 10

图 10   Box-Muller变换生成的端面一维随机噪声分布

Fig.10   One-dimensional random noise distribution of end face generated by Box-Muller transform


图 11

图 11   MH采样算法采样结果

Fig.11   Sampling results of MH sampling algorithm


图 12

图 12   采样数据分布与原始数据分布的对比图

Fig.12   Comparison between distribution of sampled data and distribution of raw data


图1112可知,MH算法采样过程曲线呈现出较好的混合度,一维和二维采样数据分布近似还原了原始数据分布,验证了所提采样方法对于一维和二维采样都能取得较好的采样效果且采样过程稳定.

将使用的采样方法构建成噪声生成模块,外接在生成式对抗网络前端,为网络生成符合原始样本分布的噪声. 当搭建WGAN-GP网络时,判别器最后一层为未激活的全连接层,设置梯度惩罚权重参数 $\lambda {\text{ = }}1$,学习率 ${\rm{lr}} = {10^{ - 4}}$,优化算法选用RMSProp.

设置SRCNN前置层学习率 ${\rm{lr}} = {10^{ - 3}}$,最后层网络学习率为 ${10^{ - 4}}$,优化算法选用Adam. 批量大小均设置为3,GPU为英伟达GTX-3080. 通过与传统GAN方法对比验证本研究所提出方法的优越性. WGAN-GP模型经过22 000个批次训练,生成的低分辨率图像作为SRCNN的输入,再对SRCNN模型训练300个批次. 本研究所提模型和对比模型生成的图像如图13所示,各模型生成器和判别器的损失函数值变化如图14所示. 图中,E为迭代次数,b为批次大小,L为损失函数值.

图 13

图 13   各模型输出结果对比图

Fig.13   Comparison diagram of output results of each model


图 14

图 14   CGAN损失函数值变化曲线图

Fig.14   Change curve of CGAN loss function value


图13可知,传统CGAN网络难以训练,在训练过程中会出现严重的梯度消失问题,导致模型经过22000个批次训练后依旧无法有效收敛. WGAN-GP相比CGAN收敛效果较好,经过7000个批次训练后,模型已经具备拟合原始图像基本轮廓的能力,但是直接向高分辨率图像建立映射关系,生成图像质量较差. 所提方法采用逐层训练的方式,避免了生成对抗网络直接向高分辨率图像映射的弊端,而是通过SRCNN提升图像分辨率,相比单独使用CGAN和WGAN-GP模型,本研究所提模型能够获得最优的生成图像质量.

图14可知,CGAN网络由于判别器过于强大,经过600次迭代后,判别器接近收敛,导致模型优化饱和,生成器无法学到有用的信息,输出图片质量较差.

图15所示为WGAN-GP损失函数变化曲线. 可以看出,随着模型迭代步数的增加,WGAN-GP的生成器和判别器损失值逐渐下降(工程实现中梯度下降法默认为最小值优化). 在迭代到约10 000次时,曲线均趋于稳定,整个训练过程没有出现梯度消失和梯度爆炸,判别器损失函数加入梯度惩罚项有效提高了训练过程的稳定性,因此WGAN-GP不需要投入大量的时间去均衡生成器和判别器的强度. SRCNN损失函数值变化如图16所示,SRCNN训练过程极其稳定,经过200次迭代后接近收敛.

图 15

图 15   WGAN-GP损失函数值变化曲线图

Fig.15   Change curve of WGAN-GP loss function value


图 16

图 16   SRCNN损失函数值变化曲线图

Fig.16   Change curve of SRCNN loss function value


各模型生成图像质量评价指标如表3所示. CGAN模型由于生成器无法学习到有效信息,生成图像质量较差,本研究所提模型相比WGAN-GP模型,生成图像的信噪比提高13.07%,平均结构相似性提高32.40%,平均FID降低37.58%.

表 3   各模型生成图像质量评价表

Tab.3  Quality evaluation table of image generated by each model

模型 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

新窗口打开| 下载CSV


为了研究不同噪声对模型训练结果的影响,分别用不同离散程度的噪声输入WGAN-GP模型进行训练,经过2 000个批次训练输出结果如图17所示,输出图片质量评价如表4所示. 由表4可知,随着输入噪声离散程度增大,生成图像的PSNR和SSIM指标均降低,FID升高,生成图像质量下降,因此,在训练时尽量采用离散程度较小的噪声.

图 17

图 17   不同标准差噪声及对应生成图像

Fig.17   Different standard deviation noise and corresponding image generation


图18为不同采样方式在WGAN-GP网络训练过程中的FID变化曲线,图中,MAP为平均精确度. 所提的预先对样本图像分布进行拟合的随机采样方法,相比直接采用均匀分布和正态分布生成噪声的方法更利于模型训练.

表 4   各模型生成图像质量评价表

Tab.4  Quality evaluation table of image generated by each model

噪声标准差 $\sigma $ 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

新窗口打开| 下载CSV


图 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

图 21   轴承滚子缺陷检测对比

Fig.21   Comparison of bearing roller defect detection


4. 结 语

对图像数据增强方法做了深入研究,提出基于生成对抗神经网络超分辨率重建的图像数据增强方法,通过案例分析和模型对比验证了所提方法具有最优的生成图像质量. 通过拟合原始数据分布构建噪声生成模块,为模型输入符合数据分布的二维噪声,有效降低了模型的参数冗余度,提高了模型生成图像质量的计算经济性. 融合WGAN-GP和SRCNN模型提出一种复合模型,避免生成对抗网络直接向高分辨率图像建立映射关系的弊端. 本研究所提出的生成对抗网络结合图像超分重建思想所用模型并不固定,下一步可以尝试不同的模型组合,例如BEGAN+SRCNN.

参考文献

康守强, 胡明武, 王玉静, 等

基于特征迁移学习的变工况下滚动轴承故障诊断方法

[J]. 中国电机工程学报, 2019, 39 (3): 764- 772

DOI:10.13334/J.0258-8013.PCSEE.180130      [本文引用: 1]

KANG Shou-qiang, HU Ming-wu, WANG Yu-jing, et al

Fault diagnosis method of a rolling bearing under variable working conditions based on feature transfer learning

[J]. Proceedings of the CSEE, 2019, 39 (3): 764- 772

DOI:10.13334/J.0258-8013.PCSEE.180130      [本文引用: 1]

肖雄, 肖宇雄, 张勇军, 等

基于二维灰度图的数据增强方法在电机轴承故障诊断的应用研究

[J]. 中国电机工程学报, 2021, 41 (2): 738- 749

DOI:10.13334/j.0258-8013.pcsee.200834      [本文引用: 1]

XIAO Xiong, XIAO Yu-xiong, ZHANG Yong-jun, et al

Research on the application of the data augmentation method based on 2D gray pixel images in the fault diagnosis of motor bearing

[J]. Proceedings of the CSEE, 2021, 41 (2): 738- 749

DOI:10.13334/j.0258-8013.pcsee.200834      [本文引用: 1]

PAN S J, YANG Q

A survey on transfer learning

[J]. IEEE Transactions on Knowledge and Data Engineering, 2009, 22 (10): 1345- 1359

[本文引用: 1]

SALAKHUTDINOV R, LAROCHELLE H. Efficient learning of deep Boltzmann machines [C]// Proceedings of the 13th International Conference on Artificial Intelligence and Statistics. Sardinia: JMLR, 2010: 693-700.

[本文引用: 1]

BENGIO Y, LAUFER E, ALAIN G, et al. Deep generative stochastic networks trainable by backprop [C]// International Conference on Machine Learning. Beijing: PMLR, 2014: 226-234.

[本文引用: 1]

KINGMA D P, WELLING M. Auto-encoding variational bayes [EB/OL]. [2013-12-20]. https://arxiv.org/abs/1312.6114.

[本文引用: 1]

VANDEN OORD A, KALCHBRENNER N, ESPEHOLT L, et al. Conditional image generation with pixel-cnn decoders [C]// 30th Conference on Neural Information Processing Systems. Barcelona: CA, 2016: 4797-4805.

[本文引用: 1]

VAN OORD A, KALCHBRENNER N, KAVUKCUOGLU K. Pixel recurrent neural networks [C]// International Conference on Machine Learning. New York: JMLR, 2016: 1747-1756.

[本文引用: 1]

GOODFELLOW I, POUGET-ABADIE J, MIRZA M, et al

Generative adversarial networks

[J]. Communications of the ACM, 2020, 63 (11): 139- 144

DOI:10.1145/3422622      [本文引用: 1]

MIRZA M, OSINDERO S. Conditional generative adversarial nets [EB/OL]. [2014-11-06]. https://arxiv.org/abs/1411.1784.

[本文引用: 1]

RADFORD A, METZ L, CHINTALA S. Unsupervised representation learning with deep convolutional generative adversarial networks [EB/OL]. [2015-11-19]. https://arxiv.org/abs/1511.06434.

[本文引用: 1]

ARJOVSKY M, CHINTALA S, BOTTOU L. Wasserstein generative adversarial networks [C]// International Conference on Machine Learning. Sydney: PMLR, 2017: 214-223.

[本文引用: 2]

GULRAJANI I, AHMED F, ARJOVSKY M, et al. Improved training of wasserstein gans [EB/OL]. [2017-03-31]. https://arxiv.org/abs/1704.00028.

[本文引用: 1]

BERTHELOT D, SCHUMM T, METZ L. Boundary equilibrium generative adversarial networks [EB/OL]. [2017-03-21]. https://arxiv.org/abs/1703.10717.

[本文引用: 1]

KARRAS T, AILA T, LAINE S, et al. Progressive growing of gans for improved quality, stability, and variation [C]// International Conference on Learning Representations. Vancouver: JMLR, 2018: 26.

[本文引用: 1]

KARRAS T, LAINE S, AILA T. A style-based generator architecture for generative adversarial networks [C]// CVF Conference on Computer Vision and Pattern Recognition. Long Beach: IEEE, 2019: 4396-4405.

[本文引用: 1]

KARRAS T, LAINE S, AITTALA M, et al. Analyzing and improving the image quality of style GAN [C]// CVF Conference on Computer Vision and Pattern Recognition. Seattle: IEEE, 2020: 8107-8116.

KARRAS T, AITTALA M, LAINE S, et al

Alias-free generative adversarial networks

[J]. Advances in Neural Information Processing Systems, 2021, 34: 852- 863

[本文引用: 1]

DONG C, LOY C C, HE K, et al

Image super-resolution using deep convolutional networks

[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2015, 38 (2): 295- 307

[本文引用: 1]

DEON A F, MENYAEV Y A

Twister generator of random normal numbers by box-muller model

[J]. International Journal of Trends in Computer Science, 2020, 16 (1): 1- 13

[本文引用: 1]

何新林, 戚宗锋, 李建勋

基于隐变量后验生成对抗网络的不平衡学习

[J]. 上海交通大学学报, 2021, 55 (5): 557- 565

DOI:10.16183/j.cnki.jsjtu.2019.264      [本文引用: 1]

HE Xin-lin, QI Zong-feng, LI Jian-xun

Unbalanced learning of generative adversarial network based on latent posterior

[J]. Journal of Shanghai Jiaotong University, 2021, 55 (5): 557- 565

DOI:10.16183/j.cnki.jsjtu.2019.264      [本文引用: 1]

ARJOVSKY M, BOTTOU L. Towards principled methods for training generative adversarial networks [EB/OL]. [2017-01-17]. https://arxiv.org/abs/1701.04862.

[本文引用: 1]

LYE A, CICIRELLO A, PATELLI E

Sampling methods for solving bayesian model updating problems: a tutorial

[J]. Mechanical Systems and Signal Processing, 2021, 159: 107760

DOI:10.1016/j.ymssp.2021.107760      [本文引用: 1]

/