FMRGAN: Feature Mapping Reconstruction GAN for Rolling Bearings Fault Diagnosis Under Limited Data Condition

Abstract

Due to the reality that it is difficult to collect sufficient and balanced data for all fault types of rolling bearings, it is a challenging mission to accurately realize the rolling bearing fault diagnosis under limited data conditions. Utilizing generative adversarial networks (GANs) to solve the limited data augmentation problem has been proven to be an effective approach. However, existing GAN-based data augmentation methods do not take into account the checkerboard artifacts caused by the use of transposed convolution in the generator, which in turn affects the quality of the generated samples. To address this issue, this article proposes a rolling bearing fault diagnosis method based on feature mapping reconstruction GAN (FMRGAN), which enhances the training sample set by generating high-quality fault data to improve the performance of the fault diagnosis model under limited data conditions. First, the vibration signals are transformed into time-frequency maps using continuous wavelet transform (CWT), and then sufficient synthetic samples are generated by FMRGAN. The feature mapping reconstruction module of the adaptive generative restructuring kernel is used to construct the generator, and the coordinate attention (CA) mechanism is introduced into the discriminator to effectively avoid the checkerboard artifacts. Second, a novel TokenDrop regularization method is designed, which contributes to the Vision Transformer classification model to capture local features in the time-frequency diagram better and reduce overfitting. Finally, the effectiveness of the proposed FMRGAN-based fault diagnosis method is validated on the CWRU dataset and the compound fault dataset collected by the self-built platform, respectively.