Full Attention Wasserstein GAN With Gradient Normalization for Fault Diagnosis Under Imbalanced Data

Abstract

The fault diagnosis of rolling bearings is vital for the safe and reliable operation of mechanical equipment. However, the imbalanced data collected from the real engineering scenario bring great challenges to the deep learning-based diagnosis methods. For this purpose, this article proposes a methodology called full attention Wasserstein generative adversarial network (WGAN) with gradient normalization (FAWGAN-GN) for data augmentation and uses a shallow 1-D convolutional neural network (CNN) to perform fault diagnosis. First, a gradient normalization (GN) is introduced into the discriminator as a model-wise constraint to make it more flexible in setting the structure of the network, which leads to a more stable and faster training process. Second, the full attention (FA) mechanism is utilized to let the generator pay more attention to learning the discriminative features of the original data and generate high-quality samples. Third, to more thoroughly and deeply evaluate the data generation performance of generative adversarial networks (GANs), a more comprehensive multiple indicator-based evaluation framework is developed to avoid the one-sidedness and superficiality of using one or two simple indicators. Based on two widely applied fault diagnosis datasets and a real rolling bearing fault diagnosis testbed, extensive comparative fault diagnosis experiments are conducted to validate the effectiveness of the proposed method. Experimental results reveal that the proposed FAWGAN-GN can effectively solve the sample imbalance problem and outperforms the state-of-the-art imbalanced fault diagnosis methods.