Data-augmented wavelet capsule generative adversarial network for rolling bearing fault diagnosis

Rolling bearing fault diagnosis with limited imbalance data is significant and challenging. It is​ a nice attempt to generate data for balancing datasets. In this paper, a wavelet capsule generative adversarial network (WCGAN) is proposed to address this issue. Firstly, the Harr wavelet is introduced into GAN to construct wavelet transform GAN (WTGAN). It keeps convolutional neural networks (CNNs) shift-invariant to extract the deep features of the data. Secondly, WCGAN is developed to alleviate CNNs’ incomplete analysis of signal information, which replaces part of CNNs in WTGAN with capsule networks. Thirdly, a novel loss function is designed for WCGAN to maintain a smooth training process and improve the quality of the generated data. Furthermore, various experiments are conducted in multiple ways to confirm the effectiveness and accuracy of the novel method. Results indicate that the proposed method balances the dataset and surpasses other advanced approaches in imbalanced data diagnosis with potential.