Imbalanced fault diagnosis of rolling bearing using improved MsR-GAN and feature enhancement-driven CapsNet

•An improved GAN-based DA strategy is proposed to extend imbalanced sample data.•Multiple similarity indexes are adopted to evaluate the quality of generated data.•A feature enhancement network is designed to mine fault-sensitive features.•A FED-CapsNet model is developed for imbalanced fault diagnosis.•The proposed approach performs well with extremely imbalanced bearing dataset. Traditional fault diagnosis approaches of rolling bearing often need abundant labeled data in advance while some certain fault data are difficult to be acquired in engineering scenarios. This imbalanced fault data problem limits the diagnostic performance. To solve it, an imbalanced fault diagnosis approach based on improved multi-scale residual generative adversarial network (GAN) and feature enhancement-driven capsule network is proposed in this paper. Firstly, frequency slicing wavelet transform is utilized to extract two-dimensional time–frequency features from original vibration signals. By designing multi-scale residual network structure and hybrid loss function, original GAN model is improved, generating high-quality fake time–frequency features to balance fault data distribution. To increase the attention of the diagnostic model to fault-sensitive features and suppress irrelevant features, a feature enhancement network is designed to dynamically weight the fault features by modeling the feature importance. On this basis, enhanced performance of imbalanced fault classification is achieved. Verification experiments demonstrate that it performs well in processing the imbalanced fault data, and has better stability and diagnostic accuracy than state-of-the-art methods.