Dual generative adversarial networks combining conditional assistance and feature enhancement for imbalanced fault diagnosis

Abstract

The dataset in the application scenario of existing fault diagnosis methods is often balanced, while the data collected under actual working conditions are often imbalanced. Directly applying existing fault diagnosis methods to this scenario will lead to poor diagnosis effect. In view of the above problems, we proposed a method called dual generative adversarial networks (DGANs) combining conditional assistance and feature enhancement. The method uses data augmentation as a basic strategy to supplement imbalanced datasets by generating high-quality data. Firstly, a new generator is designed to build the basic framework by sharing the dual-branch deconvolutional neural networks, and combining the label auxiliary information and the coral distance loss function to ensure the diversity of generated samples. Secondly, a new discriminator was designed, which is based on deep convolutional neural networks and embedded with auxiliary classifiers, further expanding the function of the discriminator. Thirdly, the self-attention module is introduced into both the generator and the discriminator to enhance deep feature learning and improve the quality of generated samples; finally, the proposed method is experimentally validated on datasets of two different testbeds. The experimental results show that the proposed method can generate fake samples with rich diversity and high quality, using these samples to supplement the imbalanced dataset, the effect of imbalanced fault diagnosis has been substantially improved. This method can be used to solve the problem of fault diagnosis in the case of sample imbalance, which often exists in actual working conditions.

Introduction

Rolling bearings are the core components of the transmission device of rotating machinery, and their reliability affects the operation of the entire rotating machinery. Therefore, timely and effective fault diagnosis of rolling bearings provides an important guarantee for the safe and reliable operation of rotating machinery.^1–3 At present, the fault diagnosis methods based on artificial intelligence have the ability to learn features independently, which avoids the drawbacks of methods based on signal processing that require a lot of prior knowledge support and human support for feature selection, achieving good diagnosis results.^4–6 So, the method based on artificial intelligence is currently the mainstream method of bearing fault diagnosis.

To ensure better and stable fault diagnosis effects of rolling bearings, methods based on artificial intelligence often require a large number of balanced samples for driving.^7–9 However, in real working conditions, the bearings of rotating mechanical equipment are often running in a normal state. Considering safety and economy, when bearing failures occur, mechanical equipment will be immediately stopped. Therefore, there are often more normal samples collected, but fewer fault samples, and the number of fault samples is far less than that of normal samples. Due to the large difference between the number of normal samples and the number of fault samples, there is an imbalance between normal samples and fault samples. It is difficult for a small number of fault samples to meet the needs of artificial intelligence fault diagnosis methods that require a large number of balanced samples to drive, resulting in the reduction of fault diagnosis accuracy.

Since the samples collected in actual working conditions are often imbalanced, imbalanced fault diagnosis has always been a difficult and hot topic in the field of fault diagnosis.^10,11 In recent years, scholars have put forward many solutions for imbalanced fault diagnosis, which can be generally summarized into three categories: data augmentation method, feature learning method, and classifier design method. ¹² The method of data augmentation starts from the perspective of data preprocessing, the original limited amount of fault data is supplemented by data generation or data oversampling. For example, Yi et al. ¹³ proposed a minority class clustering SMOTE method to improve the imbalanced classification performance. Liu et al. ¹⁴ introduced an imbalanced fault diagnosis method based on improved multi-scale residual generative adversarial network (GAN) and feature enhancement-driven capsule network, which achieved the enhanced performance of imbalanced fault classification. The method based on feature learning takes feature extraction as an entry point, without generating data, and learns fault features from limited fault data by designing a regularized neural network or feature extraction model. For instance, Saufi et al. ¹⁵ designed a deep learning model based on stacked sparse autoencoder model to deal with limited sample data, and achieved good diagnostic performance on experimental gearbox and wind turbine gearbox datasets respectively. Li et al. ¹⁶ raised a fault diagnosis method called deep balanced domain adaptation neural network for fault diagnosis, which achieved satisfactory results in the case of limited labeled data of planetary gearbox. The method based on classifier design starts from the aspect of conditional classification, and directly classifies the imbalanced samples by designing a classifier adapted to the imbalanced samples, without generating data or designing a feature extraction model. For instance, Jia et al. ¹⁷ came up with a framework called deep normalized convolutional neural network (CNN) and developed a weighted Softmax class weighting strategy to classify imbalanced faults. He et al. ¹⁸ presented a support tensor machine with dynamic penalty factor and applied it to the fault diagnosis of rotating machinery, which achieved good classification results when the training samples were imbalanced data. Although the methods based on feature learning and methods based on classifier design have achieved good diagnostic results in imbalanced fault diagnosis, they cannot be flexibly applied to the imbalanced scene with variable number of fault samples due to the limitation of their basic motivation (directly processing imbalanced data), which leads to the poor generalization ability of the model and certain limitations in the application scope. However, the method based on data augmentation avoids the drawbacks of the above methods due to different basic motivations (supplemental data to imbalanced data), and becomes a very practical and promising method to improve the accuracy of imbalanced fault diagnosis and ensure the generalization ability of the model.

GAN ¹⁹ is a data augmentation method. Comparing with other data augmentation methods, such as SMOTE, ADASYN, etc., GAN is often used for data enhancement, due to its own structural characteristics, which has the ability to continuously learn. Numerous GAN variant networks emerged. For example, Wang et al. ²⁰ proposed an enhanced GAN, based on the use of degree convolution Generative Adversarial network to generate samples, the integral K-means clustering algorithm was introduced into the CNN fault diagnosis model to improve the classification accuracy. Zhang et al. ²¹ put forward a fault diagnosis method based on deep learning to address the data imbalance problem by explicitly creating additional training data. Zhou et al. ²² used a scheme of global optimization to design a new generator and discriminator, aiming to generate those fault features extracted from a few fault samples via Auto Encoder instead of fault data sample. Zhao and Yuan ²³ raised an improved GAN to improve the fault diagnosis performance of imbalanced data. The above-mentioned GAN variant networks have achieved good diagnostic results in imbalanced fault diagnosis. However, since the basic structure adopted the GAN model, the adversarial mechanism brought by its unique structural characteristics can also make it difficult to maintain a balance between the generator and the discriminator; there are still problems of mode collapse or gradient disappearance, which may lead to insufficient diversity of generated samples and unstable convergence of model, affecting the quality of generated samples. In order to better overcome the above problems, this paper proposed dual generative adversarial networks (DGANs) combining conditional assistance and feature enhancement for imbalanced fault diagnosis. The method combines the label auxiliary information, and integrates the shared dual-branch strategy that embed with the coral distance to increase the diversity of generated samples; meanwhile, a self-attention module is introduced to enhance the depth features to improve the quality of the generated samples.

The method proposed in this paper is a novel data augmentation method, and the detailed description is as follows. Firstly, the label auxiliary information is introduced into both the generator and the discriminator to ensure the diversity of sample generated. Secondly, a new generator is designed with a shared dual-branch deconvolutional neural network (DCNN), and the coral distance loss function is embedded to increase the diversity of generated samples by maximizing the coral distance. Thirdly, a CNN is used to design a discriminator to extract useful feature information and degradation information of input samples, and an auxiliary classifier is embedded to further expand the function of the discriminator. Then, the self-attention module is introduced into the generator and the discriminator to promote the ability of the network to automatically learn important global information, improve the discrimination performance of discriminator and the ability of generator to generate samples, and compensate for the possible reduction in similarity due to sample diversity. Finally, through continuous adversarial training between the generator and the discriminator, the generator generates fake samples with rich diversity and high quality, and supplements the generated fake samples into the imbalanced dataset for fault diagnosis. The main contributions of the proposed method are as follows:

1. A new generator is designed using a shared dual-branch DCNN, and a new loss function is designed to help the generation of sample diversity by embedding label auxiliary information and a coral distance loss function to prevent mode collapse.

2. A new discriminator is designed by using CNN and embedding auxiliary classifier, and a self-attention module is introduced in both the discriminator and the generator to enhance deep feature learning and improve the quality of generated samples.

3. The experimental results show that the proposed method is used to generate data and supplement imbalanced datasets to substantially improve the performance of imbalanced fault diagnosis.

The rest of this paper is structured as follows. In section of theoretical background, the theoretical background of the basic knowledge used in the proposed method is introduced in detail. In section of the proposed method, the specific process of the proposed method is introduced in detail. In section of experimental verification, the effectiveness and superiority of the proposed method are proved by experiments on two different datasets. In section of conclusion, the proposed method is summarized and further prospected.

Theoretical background

Convolutional neural networks and deconvolutional neural networks

Convolutional neural networks

CNNs are widely used in the field of fault diagnosis.^24,25 The deep CNNs can extract the deep features of the input signal through the deep network, and can make effective judgments on the extracted deep features and identify the fault type. Deep CNNs are mainly composed of convolutional layers, activation layers, pooling layers, and fully connected layers.

The convolution layer actually uses the convolution kernel to downsample the input signal, and the result obtained after downsampling is expressed as a specific local feature extracted from the input signal. The mathematical expression of the convolution operation is as follows:

$$\mathrm{con}{\mathrm{v}}^{[l]}=W^{[l]}*x^{[l-1]}+b^{[l]}$$

(1)

where, $W^{[l]}$ represents the weight of the convolution kernel in the l convolutional layer. $b^{[l]}$ represents the bias of the convolution kernel in the l convolutional layer. $x^{[l-1]}$ represents the output of the $l-1$ convolutional layer and also represents the input of the l convolutional layer. $\mathrm{con}{\mathrm{v}}^{[l]}$ represents the output of the l convolutional layer.

The main function of the activation layer is to non-linearly map the result after the convolution operation. Common activation functions mainly include tanh, Sigmoid, Relu, LeakyRelu, Selu, etc. The corresponding activation function can be selected according to the network architecture and actual needs. For example, the mathematical expression of LeakyRelu is as follows:

$$\begin{array}{cc}\multicolumn{1}{c}{\mathrm{act}{\mathrm{i}}^{[l]}=}& \mathrm{LeakyRelu}({\mathrm{conv}}^{[l]})\hfill \\ \multicolumn{1}{c}{=\{\begin{array}{c}{\mathrm{conv}}^{[l]},\mathrm{if}{\mathrm{conv}}^{[l]}\ge 0\hfill \\ \multicolumn{1}{c}{\alpha {\mathrm{conv}}^{[l]},\mathrm{if}\mathrm{con}{\mathrm{v}}^{[l]}<0}\end{array}}\end{array}$$

(2)

where, the value of $\alpha$ is generally set to 0.01, the purpose is to prevent the gradient direction aliasing problem that occurs when the input is less than 0 during the gradient backpropagation process.

The pooling layer further filters the activated features, extracts representative features that can better represent the original signal, and can effectively reduce the scale of the output features, reduce the amounts of parameters of the model, and improve the computational efficiency of the model. The pooling layer is usually divided into maximum pooling and average pooling according to different operation types. The most commonly used is maximum pooling, which takes the largest feature value extracted in the receptive field as the output. The mathematical expression is as follows:

$${\mathrm{pool}}^{[l]}=\max \left\{{\mathrm{acti}}^{[l]}\right\}$$

(3)

The fully connected layer is responsible for recombining the pooled features and mapping the multi-dimensional feature input to the two-dimensional feature output. The mathematical formula is as follows:

$$\mathrm{f}{\mathrm{c}}^{[l+1]}=W^{[l+1]}*\mathrm{poo}{\mathrm{l}}^{[l]}+b^{[l+1]}$$

(4)

Among them, $W^{[l+1]}$ and $b^{[l+1]}$ represent the weight and bias of the convolution kernel in the $l+1$ fully connected layer respectively, and $\mathrm{f}{\mathrm{c}}^{[l+1]}$ represents the output of the $l+1$ fully connected layer.

The output of the fully connected layer is regressed or classified according to task requirements. Generally, classification tasks are most common in fault diagnosis. Softmax activation function is generally used for classification, and its mathematical expression is as follows:

$$\begin{array}{c}\multicolumn{1}{c}{\mathrm{out}=\mathrm{Soft}max\left(\mathrm{fc}\right)=\left[\begin{array}{c}p(C=1|\mathrm{fc};\theta )\hfill \\ \multicolumn{1}{c}{p(C=2|\mathrm{fc};\theta )}\\ \multicolumn{1}{c}{p(C=c|\mathrm{fc};\theta )}\end{array}\right]}\\ \multicolumn{1}{c}{=\frac{1}{\sum _{n=1}^c\exp \left({\theta }_n^T\mathrm{fc}\right)}\left[\begin{array}{c}\exp \left({\theta }_1^T\mathrm{fc}\right)\hfill \\ \multicolumn{1}{c}{\exp \left({\theta }_2^T\mathrm{fc}\right)}\\ \multicolumn{1}{c}{\exp \left({\theta }_c^T\mathrm{fc}\right)}\end{array}\right]}\end{array}$$

(5)

where ${\theta }_1,{\theta }_2,···{\theta }_C$ are the regularization parameters. c represents the category of the classification. C is the sample label. $\mathrm{out}$ represents the classified output of training data after passing through the network.

Finally, calculate the classification loss function of Softmax, the mathematical formula is as follows:

$$L_C=-\sum _i{C}_i\log (\mathrm{out})$$

(6)

Among them, $C_i$ represents the value of the i-th position in the actual label, and $L_C$ represents the final classification loss of the classifier.

Deconvolutional neural networks

DCNNs is also called transposed CNNs. As the name implies, it is generally opposite to the function of CNNs. Each layer of DCNN can be regarded as the inverse process of the corresponding layer of CNNs. ²⁶ Therefore, the DCNNs can be understood as the process of reconstructing the features extracted by the CNNs and restoring the input. The general composition of DCNNs is composed of deconvolution layers, activation layers, and unpooling layers. The deconvolution layer is a deconvolution operation on the original feature matrix to expand its dimension to match the input of the corresponding convolution layer. Its input-to-output dimensional transformation relationship is exactly the opposite of that of convolution. The operation of the activation layer is consistent with the operation of the CNNs, both to increase the nonlinearity of the network. The operation process of the unpooling layer is also the inverse process of the pooling layer of the CNNs, to supplement the reduced features and further restore the original input. For example, the unpooling operation process of maximum pooling is to place the input value at the position index of the maximum value recorded during the corresponding maximum pooling, and other positions are supplemented with 0, and the input size is equal to the output size of pooling layer of the corresponding CNNs. Due to the special properties of the DCNNs, it has a wide range of applications, such as image restoration and speech recognition.

Auxiliary classifier generative adversarial networks

GAN is a classic adversarial neural network, which is mainly composed of generator and discriminator, as shown in Figure 1(a). The core of GAN is to introduce the idea of the game in Go. The generator and the discriminator continue to play, prompting the generator to generate fake samples that are closer to the real samples. The overall optimization objective formula is as follows:

$$\begin{array}{c}\multicolumn{1}{c}{L_{\mathrm{GAN}}=E_{x~P_{\mathrm{Data}}}[\log P(\mathrm{Source}=\mathrm{real}|x_{\mathrm{real}})]}\\ \multicolumn{1}{c}{+E_{z~P_z}[\log P(\mathrm{Source}=\mathrm{fake}|x_{\mathrm{fake}})]}\\ \multicolumn{1}{c}{=E_{x~P_{\mathrm{Data}}}[\log D(x)]+E_{z~P_z}[\log (1-D(G(z)))]}\end{array}$$

(7)

where, x represents the real sample, G(z) represents the generated sample, and z represents the random noise; D(x) and D(G(z)) represent the discriminant probability of the discriminator to the real sample and the generated sample respectively; $E_{x~P_{\mathrm{data}}}$ represents the expectation of the real sample distribution, and $E_{z~P_z}$ represents the expectation of the generated sample distribution.

Figure 1.

The basic structure of (a) GAN, (b) CGAN, and (c) ACGAN.

Auxiliary classifier generative adversarial networks (ACGANs) ²⁷ is gradually developed on the basis of GAN, and its main purpose is to solve the problem of GAN mode collapse. The formation process is mainly divided into two stages. First, the category labels are added as auxiliary information on the basis of the GAN model, namely conditional generative adversarial networks (CGAN), ²⁸ as shown in Figure 1(b). CGAN overcomes the disadvantage that the original GAN cannot control the categories of generated samples, and then guides the sample generation process. The overall optimization objective function of CGAN is as follows:

$$\begin{array}{c}\multicolumn{1}{c}{L_{\mathrm{CGAN}}=E_{x~P_{\mathrm{Data}}}[\log P(\mathrm{Source}=\mathrm{real}|x_{\mathrm{real}})]}\\ \multicolumn{1}{c}{+E_{z~P_z}[\log P(\mathrm{Source}=\mathrm{fake}|x_{\mathrm{fake}})]}\\ \multicolumn{1}{c}{=E_{x~P_{\mathrm{Data}}}[\log D(x|C)]+E_{z~P_z}[\log (1-D(G(z|C)))]}\end{array}$$

(8)

Then CGAN is further improved by embedding auxiliary classifiers in the discriminator, named ACGAN, as shown in Figure 1(c). This improvement further expands the function of the discriminator, which can not only judge the source of the sample (real or fake) but also classify the sample category. So, the loss function of ACGAN is mainly composed of discriminative loss and classification loss. The mathematical expressions of the loss function are expressed respectively:

$$\begin{array}{c}\multicolumn{1}{c}{L_{\mathrm{Source}}=E_{x~P_{\mathrm{Data}}}[\log P(\mathrm{Source}=\mathrm{real}|x_{\mathrm{real}})]}\\ \multicolumn{1}{c}{+E_{z~P_z}[\log P(\mathrm{Source}=\mathrm{fake}|x_{\mathrm{fake}})]}\end{array}$$

(9)

$$L_{\mathrm{Class}}=E_{x~P_{\mathrm{Data}}}[\log P(C=c|x_{\mathrm{real}})]+E_{z~P_z}[\log P(C=c|x_{\mathrm{fake}})]$$

(10)

To achieve the purpose that both the discriminator and the classifier expect the samples generated by the generator to be the samples of the specified class, the overall optimization objective of the discriminator is to maximize $L_{\mathrm{Source}}+L_{\mathrm{Class}}$ and the overall optimization objective of the generator is to minimize $L_{\mathrm{Source}}-L_{\mathrm{Class}}$ .

Distance criterion

Kullback–Leibler divergence

Kullback–Leibler (KL) divergence ²⁹ is used to measure the difference between two distributions. The value range of KL divergence is $[0,+\infty ]$ . The smaller the KL divergence is, the smaller the difference between the two distributions, and the more similar the two distributions. Assuming that the probability distributions of two variables are P and Q respectively, the mathematical formula of the KL divergence between the two is expressed as follows:

$$\mathrm{Di}{\mathrm{s}}_{\mathrm{KL}}(P||Q)=\sum _{i=1}^{n}P(i)\log \frac{P(i)}{Q(i)}$$

(11)

Maximum mean discrepancy

Maximum Mean Discrepancy (MMD) ³⁰ is also a function used to measure the distance between two distributions, but it is different from the general distance measurement function. It maps the two distributions into the Reproducing Kernel Hilbert Space (RKHS), and then measures the distance between two distributions in the RKHS. The value range of MMD divergence is also $[0,+\infty ]$ . The smaller the MMD is, the more similar the two distributions are. When MMD is 0, it means that the two distributions are exactly the same. Assuming that the probability distributions of two variables are P and Q respectively, the distance formula of MMD is expressed as follows:

$$\mathrm{Di}{\mathrm{s}}_{\mathrm{MMD}}(P,Q)=\Vert \frac{1}{n_1}\sum _{i=1}^{n_1}\varphi (P_i)-\frac{1}{n_2}\sum _{j=1}^{n_2}\varphi (Q_j){\Vert }_{\mathrm{RKHS}}^2$$

(12)

Among them, $\varphi (.)$ represents the mapping, $n_1$ and $n_2$ represent the number of samples of the two distributions respectively, and RKHS represents that the distance is measured by $\varphi (.)$ mapping the two-domain data into RKHS.

The proposed method

Aiming at the problem of sample imbalance, this paper proposed DGANs combining conditional assistance and feature enhancement. The overall framework of the proposed method is shown in Figure 2. The main content of the proposed framework includes three aspects: the diversity of generated samples using coral, deep feature learning using self-attention, and the specific process of the proposed method.

Figure 2.

The overall framework of proposed method.

Diversity of generated samples using coral

GAN uses the idea of a game to continuously train the model to generate fake samples. But it may lead to less diversity of generated samples, making mode collapse. To solve the mode collapse problem that often occurs in GANs, we design a new generator. Firstly, the basic framework of the generator is built by using shared dual-branch DCNNs. Then adding the class label auxiliary information. Finally, the loss of the coral ³¹ is introduced into the last deconvolution layer. Here, coral embedded in the deep deconvolution network is called deep coral. Deep coral is a high-dimensional spatial feature alignment method, which completes the alignment of second-order statistics of two domain features by calculating the mean value and covariance. In addition, deep coral is also a distance measurement method and can be used as a similarity measurement method to calculate the similarity degree of two domain features in high-dimensional space. According to the requirements, the function can be minimized to increase the similarity of two domain features in high-dimensional space or the function can be maximized to reduce the similarity of two domain features in high-dimensional space to increase the diversity of two domain features. According to the proposed method requirements, maximizing the coral loss of the reconstructed features of the two in the last deconvolution layer of the shared dual-branch in high-dimensional space, the generated fake samples tend to be more evenly distributed, increasing the diversity of generated samples. The coral loss function is as in Equation (13).

$$\begin{array}{c}\multicolumn{1}{c}{L_{\mathrm{coral}}=\frac{1}{4d^2}\Vert \mathrm{Cov}{\mathrm{a}}_S-\mathrm{Cov}{\mathrm{a}}_T{\Vert }_F^2}\\ \multicolumn{1}{c}{=\frac{1}{4d^2}\Vert \frac{1}{n_S-1}(D_S^T{D}_S-\frac{1}{n_S}{(1^T{D}_S)}^T(1^T{D}_S))}\\ \multicolumn{1}{c}{-\frac{1}{n_T-1}(D_T^T{D}_T-\frac{1}{n_T}{(1^T{D}_T)}^T(1^T{D}_T)){\Vert }_F^2}\end{array}$$

(13)

where, $\mathrm{Cov}{\mathrm{a}}_S$ and $\mathrm{Cov}{\mathrm{a}}_T$ represent the source domain and target domain feature covariance matrix respectively. $D_S$ and $D_T$ represent source domain data and target domain data respectively. $n_S$ and $n_T$ represent the number of samples in the source and target domains respectively. d represents the number of outputs of neurons in the layer where the coral is located. 1 is a column vector where all the entries are equal to 1. F stands for the Frobenius norm, which is to take the sum of squares of the corresponding elements and then take the result root.

In order to increase the diversity of generated samples, the coral distance loss function between two shared double branches is maximized. In addition, without adding additional computing resources, the new generator can also generate diverse fake samples. To achieve the above objectives, in the part of the generator, we only use the first branch to generate samples to complete the competition with the new discriminator and thus generate high-quality diverse fake samples. Therefore, the objective loss function of the new generator is as Equation (14):

$$\begin{array}{c}\multicolumn{1}{c}{L_G=E_{z_1~P_{z_1},c~P_c}[\log (1-D(G(z_1,c)))]+L_C(c|G(z_1,c))}\\ \multicolumn{1}{c}{-L_{\mathrm{coral}}[(c|G_f(z_1,c)),(c|G_f(z_2,c))]}\end{array}$$

(14)

where, $z_1$ and $z_2$ represent the random noise that input to the shared dual-branch generator respectively. c represents label auxiliary information. E represents mathematical expectations. $G(z_1,c)$ stands for the generated fake sample. $D(G(z_1,c)$ represents the probability that the discriminator discriminates the generated samples as fake samples. $G_f(z,c)$ represents the feature restored by the generator in the embedded layer of the coral.

In order to generate fake samples that are closer to the real sample distribution and ensure the diversity of generated samples, the final optimization objective of the generator is to minimize the generator objective loss function.

Deep feature learning using self-attention

In the field of fault diagnosis, convolutional networks, recurrent networks, ³² or fully connected networks ³³ are usually used to build diagnostic models. However, due to the limitation of the local receptive field, the convolutional network only models the local dependencies of the input information, and only has long-distance dependencies when it has a multi-layer convolutional network, but increasing the number of layers will reduce the computational efficiency of the network. Due to the capacity of information transmission and the problem of gradient disappearance, the recurrent network is actually a local encoding method, and it can only establish short-distance dependencies. Although the fully connected network establishes long-distance dependencies between input sequences, its parameters are too large, which may cause it to fail to calculate when dealing with long input sequences.

The attention mechanism ³⁴ is a computer imitation of the internal process of biological observation behavior. Its essence is to align internal experience with external sensation to increase the observation precision of some areas. Therefore, the attention mechanism can quickly capture the important information of sparse data, and focus on important information to extract its features. The self-attention mechanism ³⁵ further improves the attention mechanism, further increasing the ability to capture the internal correlation of data or features by reducing the dependence on external perception. So, a self-attention module is introduced to establish long-range dependencies between input sequences, enhance the feature learning ability, realize the learning of the deep features of the original vibration signal, endow the generator and the discriminator with stronger global important information learning ability, and improve the discriminator’s discriminative function, while promoting the generator to generate higher quality fake samples.

The structure of the self-attention mechanism is shown in Figure 2. The detailed process of the whole module is as follows:

• Step 1: The convolutional feature maps extracted by the CNNs after three 1X1 convolution operations with different number of output channels and constant size of feature map, the three 1X1 convolution operations are denoted as f(x), g(x) and h(x) respectively.

• Step 2: Transposing the output of f(x), and multiply the transposed result with the output of g(x), the output result is $s_{\mathrm{ij}}$ , and the mathematical formula is expressed as:

$$s_{\mathrm{ij}}=f(x_i{)}^{T}g(x_j)$$

(15)

• Step 3: Performing Softmax normalization on $s_{\mathrm{ij}}$ to get attention map ${\alpha }_{i,j}$ , the mathematical expression is as follows:

$${\alpha }_{i,j}=\frac{\exp (s_{\mathrm{ij}})}{\sum _{i=1}^N\exp (s_{\mathrm{ij}})}$$

(16)

where, ${\alpha }_{i,j}$ represents how much the model pays attention to the i-th position when synthesizing the j-th region.

• Step 4: Multiplying the attention map ${\alpha }_{i,j}$ with the output of h(x) to obtain the self-attention feature maps, the mathematical expression is as follows:

$$o_i=\sum _{i=1}^N{\alpha }_{i,j}h(x_i)$$

(17)

• Step 5: The self-attention feature map $o_i$ is output, and the final output feature map is y. The formula is as follows:

$$y=\gamma o_i+x_i$$

(18)

where, $\gamma$ is the scale parameter, which is initialized to 0. Let the network learn local information first, and then slowly learn long-distance information.

The specific process of the DGANs

This paper proposed DGANs that combine conditional assistance and feature enhancement to solve the problem of sample imbalance. The specific implementation processes are as follows:

• Step 1: Data acquisition. Using the vibration acceleration sensor and combining with the data acquisition system to collect the vibration signal of the bearing in the gearbox.

• Step 2: Data preprocessing. Dividing the continuous vibration signal through a sliding window of a certain length to obtain partial samples. Then, the obtained samples are subjected to fast Fourier transform (FFT) to obtain frequency domain samples to obtain more useful information to represent rolling bearing faults. Finally, the Normalizer function encapsulated in python is used for normalization, so that the values of each dimension of the sample are kept within a certain range to avoid large deviations.

• Step 3: Data generating. Introducing label auxiliary information and embedding the coral loss function in the generator to improve the diversity of the samples generated by the generator and prevent the mode collapse. Adding an auxiliary classifier to the discriminator further expands the function of the discriminator. Self-attention modules are introduced in both the discriminator and the generator to improve the quality of the generated samples. By training the generator and the discriminator to continuously confront each other, when the training is completed, the model generates high-quality fake samples.

• Step 4: Fault diagnosis. The original imbalanced dataset is supplemented with samples generated by the proposed method, and the original imbalanced state is changed by continuously supplementing the number of samples to reduce the imbalanced ratio. The supplemented dataset is used as the training dataset, and the dataset composed of the original real samples is used as the testing dataset. Considering that CNNs can extract useful fault information from the feature matrix composed of time and frequency domain features, CNNs is used as a fault diagnosis model. The detailed configuration of the CNNs is shown in Table 1. The training dataset is input into the CNNs model for training. When the training is over, the testing dataset is input into the CNNs model for testing, and the test results reflect the ability of the proposed model to generate fake samples and the effect of the proposed method to solve the samples imbalanced problem.

Table 1.

The detailed configuration of the CNNs for fault diagnosis.

Layer type	Kernel size	Output channel	Pooling	Activation function	Batch normalization
Convolution1	16*1	16	Max pooling	Relu	Yes
Convolution2	3*1	32	Max pooling	Relu	Yes
Convolution3	3*1	64	Max pooling	Relu	Yes
Convolution4	3*1	128	Max pooling	Relu	Yes
Convolution5	1*1	256	Max pooling	Relu	Yes
Full connection layer	1792*100	–	–	Relu	Yes
Full connection layer	100*4	–	–	Softmax	–

CNN: Convolutional neural network.

Among them, the new discriminator uses deep CNNs architecture, as shown in Figure 3. At the same time, the real samples with a dimension of 600 and the fake samples generated by the generator are input into discriminator, and 4 convolution layers are used to extract the signal deep features, and the extracted features are used for real and fake discrimination and fault classification. The new generator is a shared dual-branch DCNNs structure, as shown in Figure 4. Each branch inputs a Gaussian random noise z with dimension 100, mean 0, and standard deviation 1, z goes through 4 deconvolution layers and 1 convolution layer to reconstruct the noise into fake samples with the same dimension as the original real samples. In order to speed up the training convergence, InstanceNorm is added to both the generator and the discriminator. The batch size for adversarial training is set to 64 and the number of iteration steps is 2000. Taking into account the adversarial balance, the learning rates of the generator and discriminator are set to 0.0002 and 0.0004 respectively.

Figure 3.

Specific structure and parameter configuration of the new discriminator.

Figure 4.

Specific structure and parameters configuration of the new generator.

Experimental verification

Datasets description

Case Western Reserve University(CWRU) Bearing Dataset, published by Case Western Reserve University, the experimental platform is shown in Figure 5. The dataset mainly includes bearing fault data at the drive end and the fan end. The drive end bearing fault data also includes 12 k drive end bearing fault data and 48 k drive end bearing fault data. The driving end bearing type is deep groove ball, the model number is SKF6205. Bearing faults are processed by Electrical discharge machining (EDM), including fault in roller, fault in inner race and fault in outer race. Each fault type contains three different diameters of damage: 0.18, 0.36 and 0.54 mm. Installing the faulty bearing into the test motor, and recording the vibration data of the motor load at 0–3 hp respectively. In this experiment, the motor load is 3 hp, the damage diameter is 0.18 mm, and the fault data of the 12 k drive end bearing is selected. For the convenience of description, the normal data is denoted as N, and fault in roller, fault in inner race and fault in outer race fault data are denoted as RF, IF, and OF in turn.

Figure 5.

The bearing dataset of Case Western Reserve University.

Private Platform Bearing dataset, Collected by the customized platform of our research group, and the experimental platform is shown in Figure 6. The rated power of the motor is 0.75 kW, and the bearing type is roller bearing, the model number is NU205EM. The bearing fault is also processed by EDM. The fault types are fault in roller, fault in inner race and fault in outer race. The damage diameter of each fault type is 0.5 mm. Installing the faulty bearing into the test motor, and recording the vibration data of the motor with speed at 1000–1300 rpm. In this experiment, the bearing fault data are selected when the speed is 1300 rpm. Marking the normal, fault in roller, fault in inner race, and fault in outer race as N, RF, IF and OF in turn.

Figure 6.

The bearing dataset of private platform.

Data preprocessing, for the datasets collected from the above two experimental platforms, to prevent information leakage, the sliding pane overlap is used to sample the data under each health state. After sampling, 400 original time domain samples are obtained in each healthy state, each sample contains 1200 data points, and 400 frequency domain samples are obtained after fast FFT, each sample contains 600 data points. Then normalize the transformed data. For normal samples, the first 200 samples are used as the training dataset, and the last 200 samples are used as the test dataset. For the fault samples, in order to simulate the actual situation, it is assumed that the number of samples is limited, so only the first 10 samples are used as the training dataset, and the last 200 samples are used as the test dataset. Details are shown in Table 2. In this experiment, 200 normal samples and 10 fault samples of different fault types were used to generate 64 samples corresponding to their health condition.

Table 2.

The specific configuration of CWRU imbalanced dataset and private imbalanced dataset used in the experiment.

Dataset	Training dataset				Testing dataset
Health condition	N	RF	IF	OF	N	RF	IF	OF
Sample number	200	10	10	10	200	200	200	200
Sample length	600	600	600	600	600	600	600	600
Label	C0	C1	C2	C3	C0	C1	C2	C3

N: normal data; RF: fault in roller fault data; IF: fault in inner race fault data; OF: fault in outer race fault data; CWRU: Case Western Reserve University.

The effectiveness of coral

The core purpose of embedding the coral in this paper is to assist in generating more diverse samples and prevent the mode collapse of the GAN. To prove the effectiveness of the coral, the samples generated by the three models of GAN, ACGAN, and ACGAN-coral (adding the coral module on the basis of ACGAN) were visualized with the original true samples by t-distributed stochastic neighbor embedding (t-SNE). On two different datasets, the t-SNE visualization results are shown in Figures 7 and 8. In the figure, circles represent imbalanced real samples and triangles represent generated samples. From the comparison of Figure 7(a) to (c), it can be seen that on the CWRU dataset, the samples generated by the ACGAN model are intra-class relatively scattered compared to the samples generated by the GAN. While compared with the ACGAN model, the clustering of samples generated by the ACGAN-coral model is relatively weak. Especially as shown by the yellow triangle. Therefore, it can be inferred that the diversity of samples generated by the model is enhanced after embedding the coral.

Figure 7.

t-SNE visualization of (a) GAN, (b) ACGAN, and (c) ACGAN-coral in CWRU dataset.

Figure 8.

t-SNE visualization of (a) GAN, (b) ACGAN, and (c) ACGAN-coral in private dataset.

In addition, from the analysis of Figure 8(a) to (c), on the private dataset, it can be seen that the ACGAN-coral model has the weakest intra-class clustered compared with the samples generated by the GAN and ACGAN models, and the ACGAN model is slightly second, the GAN model is the best. Especially as shown by the green triangle in the picture. Therefore, on the private dataset, it is also proved that the coral can assist the model to generate more diversity of samples, and prevent mode collapse of GAN to a certain degree.

The effectiveness of self-attention

This section verifies the effectiveness of self-attention. In order to make the conclusions more accurate, two network architectures are set up here, namely ACGAN and ACGAN-attention (self-attention module added on the basis of ACGAN). Under the three fault types of two datasets, the original real sample, the sample generated by ACGAN, and the sample generated by ACGAN-attention are visualized by spectrogram. The experimental results of the two datasets are shown in Figures 9 and 10 respectively. As can be seen from Figure 9, on the CWRU dataset, compared with ACGAN, the frequency domain signal of ACGAN-attention with self-attention module is significantly enhanced at the signal shock compared with the original real sample, especially marked by the red square in the Figure 9. While the frequency domain signals generated by the ACGAN model are roughly the same at the signal shock compared with the original real sample, and even individual fault categories are weaker than the original real sample, such as RF.

Figure 9.

Spectrogram of (a) Real samples, (b) ACGAN generated samples, and (c) ACGAN- attention generated samples in CWRU dataset.

Figure 10.

Spectrogram of (a) Real samples, (b) ACGAN generated samples, and (c) ACGAN- attention generated samples in private dataset.

In addition, on the private dataset, it can be concluded from Figure 10 that under various fault types, at the signal shock, the sample in frequency domain generated by the ACGAN-attention model has also been significantly enhanced compared to the original real sample, as shown in Figure 10 marked by the red square. The frequency domain signal of the samples generated using ACGAN remains roughly the same as the original real samples at the shock signal, and even weaker than the original real samples under the IF fault type. Therefore, through the analysis of Figures 9 and 10, it can be comprehensively inferred that the use of self-attention can enhance the features to a certain extent.

Sample generation and quality assessment

The core idea of this paper is to use only a small number of fault samples to generate fake samples with similar distribution to real fault samples through adversarial training. Therefore, before supplementing the imbalanced dataset with fake samples, the quality of the generated fake samples needs to be evaluated to measure whether the generated fake samples can be used to supplement the imbalanced dataset. To prevent the accidental occurrence of a single quality assessment method, the quality of the generated samples is assessed by using two statistical indicators with large differences in the principles of MMD and KL divergence. In addition, it is compared with other methods. The MMD and KL divergence results of each method on two different datasets are shown in Table 3.

Table 3.

Details of training set and testing set.

Dataset		CWRU dataset			Private dataset
Fault type		RF	IF	OF	RF	IF	OF
MMD	GAN	5.0131	2.786	3.5844	3.8181	3.1726	3.5029
	ACGAN	1.2718	2.3896	1.8477	0.5050	0.5789	0.7954
	ACGAN-attention	1.0376	2.2636	1.6682	0.4861	0.5288	0.7207
	ACGAN-coral	1.5417	3.0853	2.3298	0.4406	0.6382	0.8026
	Proposed	0.7921	1.5767	1.1646	0.4869	0.4846	0.6849
KL	GAN	0.025	0.0096	0.0095	0.045	0.0295	0.0203
	ACGAN	0.0066	0.0108	0.0073	0.0089	0.0089	0.0082
	ACGAN-attention	0.0062	0.0099	0.0069	0.0086	0.0087	0.008
	ACGAN-coral	0.0074	0.0137	0.0088	0.0086	0.0089	0.0079
	Proposed	0.0062	0.0093	0.0071	0.0083	0.0083	0.0072

ACGAN: Auxiliary classifier generative adversarial networks; CWRU: Case Western Reserve University; GAN: generative adversarial network; MMD: Maximum Mean Discrepancy; KL: Kullback–Leibler; RF: fault in roller fault data; IF: fault in inner race fault data; OF: fault in outer race fault data.

According to the analysis in Table 3, whether on the CWRU dataset or on the private dataset, compared with other methods, the MMD and KL divergence of the proposed method under various fault types are relatively small. So, the results indicate that the data generated using the proposed method has a high similarity to the real data distribution. The MMD and KL divergence of ACGAN-coral are larger than those of ACGAN, which indicates that the use of coral increases the diversity while reducing the similarity. While adding a self-attention module on the basis of ACGAN-coral, that is, the proposed method, MMD and KL divergence are greatly reduced, and the distribution of generated samples is closer to the original real sample distribution. It is further proved that self-attention can make up for the reduction of similarity caused by sample diversity and improve the quality of generated samples.

Imbalanced fault diagnosis

At present, fault diagnosis methods can be roughly divided into two categories: one is based on signal processing, and the other is based on deep learning. The fault diagnosis method based on signal processing is to judge the fault category of the signal according to the specific fault characteristic frequency, and the fault discrimination according to the specific fault characteristic frequency needs to match the prior diagnosis knowledge of experts. The fault diagnosis method based on deep learning can identify the fault characteristics mainly based on mathematical characteristics and statistical characteristics.^36,37 The fault characteristics are not specific fault characteristic frequency corresponding to the signal processing but the edge frequency oscillation peak fault characteristic information excited by the fault characteristic frequency in the full frequency band (low frequency band, medium frequency band, high frequency band), and these fault characteristic information represents the difference between the signals of fault types. The fault diagnosis method based on deep learning is to further mine and extract fault characteristic information through the deep network, and send the extracted characteristic into the classifier to automatically complete the fault identification, so as to realize the fault diagnosis. Therefore, the fault diagnosis method based on deep learning is to automatically discriminate faults without the prior knowledge of experts. Therefore, at present, the fault diagnosis method based on deep learning is increasingly favored by scholars in the fault diagnosis field, and has also become the mainstream method in the current fault diagnosis field.

The fault diagnosis method of CNNs set earlier in this paper is one of the deep learning fault diagnosis methods, it is mainly a deep learning fault diagnosis method similar to image processing, image classification. It also classifies these images like image classification. As long as these image categories are distinguished, the convolution deep learning fault diagnosis method can distinguish the fault categories. As mentioned in above sections, the synthetic data generated by the data generation network model (DGANs) actually contains the characteristics of the fault frequency information, and the fault frequency information is not the fault characteristic frequency in the traditional sense, but the edge frequency oscillation peak fault characteristic information excited by the fault characteristic frequency in the full frequency band (low frequency band, medium frequency band, high frequency band). Therefore, the fault diagnosis method of CNNs can correctly distinguish the fault category of the signal according to the fault characteristic information image generated by the signal in the full frequency band. It is also the theoretical basis for improving the accuracy of imbalanced fault diagnosis by supplementing the synthetic data generated by the DGANs. The experiment confirmed as follows:

From the experiments in the previous sections, it is proved that the proposed method has outstanding advantages compared with other methods, the quality of the generated samples is higher, and the diversity is rich. To further verify the effectiveness of the proposed method for imbalanced fault diagnosis, in this section, for different fault types, on two different datasets, we compose a new training dataset by adding different number of generated samples to the original imbalanced training datasets, and the number of added samples for each fault type is set to 10, 20, 30, and 40 in turn.

The fault diagnosis model CNNs is trained by different training datasets that were set earlier. When the model training is stable, the test datasets that were set in the early stage is used to test the trained model. The average test results of seven consecutive times on two different datasets are shown in Table 4. From Table 4, it can be concluded that in the CWRU datasets, with the increasing number of supplementary samples, the test accuracy basically maintains a gradual upward trend. Especially when the number of supplementary samples is 10, the average test accuracy is 30.18% higher than the original imbalanced dataset, reaching 99.68%. When the number of supplementary samples is 30, the average test accuracy has stabilized at 100%. On private datasets, the test accuracy keeps increasing after using the proposed method to supplement the original imbalanced dataset with different numbers of samples. When the number of supplementary samples is 40, the average test accuracy increases from the original 57.93% to 91.65%, the average increase is 33.72%.

Table 4.

Details of training set and testing set.

Datasets		Training datasets (number of supplementary samples for each fault type)
		0	10	20	30	40
Average test accuracy	CRWU	69.50%	99.68%	99.93%	100%	100%
	Private	57.93%	75.04%	84.84%	89.29%	91.65%

CWRU: Case Western Reserve University.

To illustrate the effectiveness of the proposed method for imbalanced fault diagnosis in more detail, on different training datasets, the test results of seven consecutive times are plotted as a column chart. On two different datasets, the column charts are shown in Figures 11 and 12 respectively. According to the analysis in Figure 11, on the dataset of CWRU, the original imbalanced dataset has a large fluctuation in the results of seven consecutive times, with a standard deviation of 5.88%. With the continuous supplement samples, the stability of test results is gradually enhanced. When the sample is replenished to 30, the test results for seven consecutive times remain at 100%. On the private dataset, the test result in seven consecutive times of original imbalanced dataset has the most prominent fluctuations with a standard deviation of 2.24%. With the increasing number of supplementary samples, the stability of the test results is gradually improving. When the number of samples is supplemented to 40, the standard deviation is reduced from 2.24% to 0.78%, and the test results for seven consecutive times are roughly maintained at about 91.65%.

Figure 11.

The testing accuracy in seven consecutive times of CWRU dataset.

Figure 12.

The testing accuracy in seven consecutive times of private dataset.

Through the above comprehensive analysis, it can be concluded that the fault diagnosis accuracy is not only greatly improved, but also diagnosis stability is guaranteed to a certain extent after using the proposed method to supplement the original imbalanced dataset. Therefore, the proposed method can effectively solve the problem of sample imbalance.

In order to more intuitively illustrate the classification effect of the proposed method on imbalanced fault diagnosis, on the CWRU dataset, the confusion matrix is drawn for the test accuracy respectively, which are the original imbalanced dataset as the training dataset and the dataset supplemented with 30 samples for each type of fault as the training dataset. The confusion matrix is shown in Figure 13. It can be seen from the analysis in Figure 13 that when the number of supplementary samples is 30, for the fault in roller, the original number of correct diagnosis samples increased from 107 to 200. For fault in inner race, it is increased from 147 to 200. For fault in outer race, it is also reduced from 138 to 200. According to the calculation, after using the proposed method to supplement the samples, the original test accuracy increased from 74% to 100%, and increased by 26%.

Figure 13.

The confusion matrix of test accuracy on CWRU dataset. (a) Imbalanced, and (b) Add 30 samples.

On the private dataset, the test accuracy of the dataset when the training dataset is the original dataset and the dataset supplemented with 40 samples for each type fault visualized using the confusion matrix respectively, and the visualization results are shown in Figure 14. It can be seen from Figure 14 that when 40 samples are added, for the fault in roller, the original 178 misclassified samples are reduced to 62. For fault in inner race, it is reduced from 152 to 5. For fault in outer race, it is also reduced from 13 to 0. Therefore, it can be calculated that after the samples are supplemented through the proposed method, the overall sample misclassification rate of the testing dataset is reduced from the original 42.875% to 8.375%, that is, the test accuracy is increased by 34.5%.

Figure 14.

The confusion matrix of test accuracy on private dataset. (a) Imbalanced and (b) Add 40 samples.

Conclusion

In this paper, DGANs combining conditional assistance and feature enhancement for imbalanced fault diagnosis is proposed. This method flexibly combines label auxiliary information, coral distance loss function and self-attention module with an adversarial model. It overcomes the problem of mode collapse or gradient disappearance that often occurs in adversarial training on GAN, and stably generates fake samples with high quality and rich diversity. The method is verified by the public bearing signal of the CWRU and the bearings signal of the customized test bench. The experimental results show that the proposed method can generate high-quality fake samples and greatly improve the accuracy of imbalanced fault diagnosis, which is helpful to detect faults using fewer imbalanced samples in actual working conditions, laying a foundation for advancing the practical application of rotating machinery health assessment.

Although the proposed method has achieved good results in imbalanced fault diagnosis, there are still many aspects that can be improved. In future studies, we will try to focus on improving two aspects. The first is the tradeoff between sample diversity and sample similarity to find the balance point that can best guarantee the needs of both. Secondly, the knowledge learned from a small number of existing samples and generated samples is transferred to the target domain where there is a certain gap in distribution, and the fault types of unknown target domain samples are accurately predicted.