Data imbalance fault diagnosis method based on improved GAN

Abstract

The invention relates to a data imbalance fault diagnosis method based on an improved GAN in the technical field of industrial equipment fault diagnosis. The method comprises the following steps: S1, data preparation and model improvement; s2, constructing a fault diagnosis model; and S3, experimental simulation verification. According to the method, original data distribution can be adaptively learned, new samples are generated to increase minority classes, and then an enhanced data set is applied to network training of a support vector machine, a multi-layer perceptron, a convolutional neural network, a residual neural network, a sparse auto-encoder and the like so as to carry out subsequent fault diagnosis.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of industrial equipment fault diagnosis, and particularly relates to a data imbalance fault diagnosis method based on an improved GAN. BACKGROUND

[0002] Rolling bearings, as the core components of rotating machinery, are widely used in the industrial fields of wind power, rail transportation, heavy manufacturing, etc., and their running state directly affects the safety and production efficiency of equipment. According to statistics, more than 50% of mechanical failures are caused by bearing failures, of which the failure of inner and outer rings accounts for up to 80%. In actual industrial scenarios, bearings are often subjected to harsh working conditions such as high speed, heavy load and strong vibration, which can easily cause fatigue spalling, cracks, wear and other faults. If not diagnosed in time, it may cause economic losses due to equipment downtime, or even safety accidents. Therefore, developing efficient and reliable bearing fault diagnosis technology is of great significance to ensure the safety of industrial production. Traditional bearing fault diagnosis mainly relies on vibration analysis, temperature monitoring, oil analysis and other technologies. The vibration analysis method collects bearing seat vibration signals for time-frequency domain analysis, although it is widely used, but requires professional experience for feature extraction; the temperature monitoring method is sensitive to thermal faults such as poor lubrication, but it is difficult to detect early-stage pitting and other local damage; the oil analysis method is only suitable for oil-lubricated bearings, and is easily disturbed by other metal particles. In recent years, intelligent diagnosis methods based on machine learning have emerged, such as support vector machines and random forests, which can realize fault classification by extracting vibration signal features. However, these methods have significant limitations: first, they are heavily dependent on balanced data sets, while fault samples are scarce in industrial sites (the ratio of normal samples to fault samples often reaches 100:1), which leads to model bias towards the majority class; second, they are sensitive to noise, and electromagnetic interference, mechanical impact and other noise in actual working conditions can significantly reduce the feature extraction accuracy; third, they have insufficient generalization ability, and when the working conditions such as equipment load and speed change, the model diagnosis accuracy decreases significantly. These problems seriously restrict the application effectiveness of fault diagnosis technology in complex industrial environments. SUMMARY

[0003] In engineering practice, the fault duration of rotating machinery is usually much shorter than its normal operation period. Considering the high labor cost, it is often not feasible to collect large-scale and long-term fault data. This directly leads to the phenomenon of data imbalance that is common in the field of fault diagnosis - the number of samples representing normal state is significantly more than fault samples, which seriously restricts the accuracy and stability of the diagnosis model.

[0004] In view of the fault diagnosis problem of rolling bearings under the condition of class imbalance, the application proposes an improved model: Wasserstein distance auxiliary classification generative adversarial network with gradient penalty. The model makes two key improvements on the basis of the original architecture: one is to optimize the loss function design, and the other is to introduce a gradient penalty term. These improvements effectively alleviate the common gradient dispersion problem in the model training process, enabling the generative network to stably learn and generate pseudo samples of each class (especially the minority class). To ensure the quality and applicability of the generated pseudo samples, the research conducts multi-dimensional similarity verification on the generated pseudo samples. Only those pseudo samples that meet the preset filling rules will be used to augment the original imbalanced dataset. Subsequently, a variety of diagnostic models are selected for verification on the augmented dataset. The simulation experiment results show that the improved GAN network not only successfully overcomes the gradient disappearance and mode collapse problems in traditional network training, but also generates high-quality pseudo samples. Moreover, the dataset augmented by the model can effectively support the normal training of diagnostic models, ultimately solving the performance degradation problem caused by sample imbalance in bearing fault diagnosis.

[0005] Industrial field bearing fault samples are extremely scarce, with a normal sample to fault sample ratio often reaching 100:1, and the fault type distribution is uneven. Traditional machine learning models rely on balanced dataset training, and when faced with such imbalanced data, the model will overfit the majority class, resulting in a sharp decline in minority class recognition rate. For example, in a dataset of 99% normal samples and 1% fault samples, even if the model judges all samples as normal, the accuracy rate is still 99%, but the actual fault miss rate is close to 100%, which cannot meet the industrial early warning needs. Existing data augmentation methods (such as SMOTE oversampling) generate samples by copying or interpolation, which easily introduces noise or causes overfitting; undersampling directly discards majority class samples, causing loss of key feature information. Therefore, there is an urgent need for a technology that can generate high-quality and diversified fault samples to fundamentally solve the data distribution imbalance problem.

[0006] The application aims to propose an improved GAN (adversarial network) based data imbalance fault diagnosis method, which can adaptively learn the original data distribution, generate new samples to increase the minority class, and then apply the augmented dataset to support vector machines (SVM), multilayer perceptron (BPNN), and convolutional neural networks (LeNet5), residual neural networks (ResNet18), and sparse autoencoders (SAE) for training to perform subsequent fault diagnosis.

[0007] To solve the above problems, the application provides an improved GAN based data imbalance fault diagnosis method, including the following steps: S1 Data preparation and model improvement: Collect the original unbalanced dataset, generate the generative adversarial network, generate a balanced data generation model capable of generating high-quality pseudo samples by alternately training the generator and the discriminator, and improve the generative adversarial network; S2 Fault diagnosis model construction: Use the improved generative adversarial network to generate high-quality pseudo samples that approximate real samples, and then input the pseudo samples and real samples into the convolutional neural network for fault diagnosis; S3 Experimental simulation verification: Add the generated pseudo samples to the real rolling bearing data step by step, and then construct a balanced dataset for subsequent fault diagnosis.

[0008] Further, step S1 specifically includes the following steps: S1.1 Construct a basic generative adversarial network: build a basic GAN model composed of a generator and a discriminator; wherein the generator is responsible for mapping random noise to generated samples; the discriminator is responsible for distinguishing whether the input sample comes from real data or the generator, and outputs the probability that it determines to be a real sample; S1.2 Introduce Wasserstein distance and gradient penalty: Introduce Wasserstein distance as a loss term to measure the difference in data distribution, and introduce a gradient penalty term in the loss function of the discriminator to meet the Lipschitz continuity condition required for Wasserstein distance calculation; S1.3 Improve the GAN model: Construct a robust generative adversarial network with Wasserstein distance as the loss and supplemented by a gradient penalty mechanism, which is used for subsequent generation of high-quality pseudo samples.

[0009] Further, in step S1.1, the random noise is sampled from the prior distribution , the real sample is sampled from the real data distribution , , the generated sample is used to train the discriminator to maximize the discriminant loss, and the generator is trained to minimize the loss, and the objective function is defined as (1).

[0010] Further, in step S1.2, the calculation formula of the Wasserstein distance as the loss term is as follows: (2) The Lipschitz condition stipulates that there is a continuous function that ensures that there is a constant K greater than zero in the entire function domain, and any two values , satisfy: (3) wherein denotes the probability distribution of real input data of real samples, denotes the probability distribution of generated pseudo data of generated samples, denotes mathematical expectation, denotes the distribution function of real input data or generated pseudo data; Gradient penalty is selected to replace the network parameter restriction in a fixed range, so as to realize the Lipschitz condition and avoid the gradient disappearance caused by improper parameter range setting. The gradient penalty formula is (4) (5) wherein denotes the cost function of gradient penalty, denotes the penalty coefficient, denotes the linear interpolation calculation on the data of the mixed sample set, , denotes uniform distribution, denotes the probability distribution between real input data and generated pseudo data, denotes derivation on is the output after the network model discriminator, is the 2-norm.

[0011] Further, in step S1.3, the Wasserstein distance using the input real data and the generated pseudo data distribution is improved to be used as a loss term, replacing the original calculation of the binary cross-entropy function part, and the model loss function is L s and L c may be respectively expressed as (6) (7) wherein denotes the probability distribution of real input data, denotes the probability distribution of generated pseudo data, denotes the probability distribution of real label data, denotes the probability distribution of predicted label data; The improved cost function of the adversarial network is the cost function of the gradient penalty added to the cost function of the original adversarial network, and the improved cost function is (8) ​(9).

[0012] Further, step S2 specifically comprises the following steps: S2.1 randomly sampling in a Gaussian distribution to generate random noise , input the random noise and the category label c into the generator to obtain generated data; S2.2 input the generated data and the real bearing data in batches into the discriminant network, the discriminant network calculates the model loss based on the real label, and optimizes the parameters through back propagation; S2.3 after completing the discriminant network training once, the generated network can be trained multiple times, the discriminant network parameters are frozen so that they are not updated with iterations, and then gradient calculation is performed based on the cost function, and the generated network parameters are optimized along the gradient direction; S2.4 repeat steps S2.1-S2.3, and dynamically adjust the training rounds of the generated network according to the loss value fluctuation; when the model reaches Nash dynamic equilibrium, i.e. the pseudo samples generated by the generator can deceive the discriminator, and the discriminator cannot accurately determine the data source but can still correctly classify the bearing fault categories, the model training is completed.

[0013] Compared with the prior art, the present application has the following technical advantages and beneficial effects: (1) In order to alleviate the extreme class imbalance of bearing fault data, the Wasserstein distance is introduced to measure the difference between the real distribution and the generated distribution, and the convolutional neural network is used as the backbone architecture, which significantly improves the training stability; further, gradient penalty (GP) is used instead of traditional weight clipping, so that the network implicitly satisfies the Lipschitz constraint, thereby eliminating gradient disappearance and mode collapse; (2) In the training stage, the label information is fully mined, and an auxiliary output branch is added, which not only ensures the high fidelity of the generated samples, but also strengthens the joint judgment ability of the discriminator on the authenticity and category attribution of the input samples, and finally realizes the synchronous leap of high-quality data augmentation and rolling bearing fault diagnosis precision. BRIEF DESCRIPTION OF DRAWINGS

[0014] Figure 1 is the flowchart of the data imbalance fault diagnosis method based on the improved GAN of the application; Figure 2 is a model structure diagram of the adversarial network of the application; Figure 3 is a flowchart of the fault diagnosis model construction of the application; Figure 4 is a model generator and discriminator loss curve diagram of the application; Figure 5 is a model classification and discrimination accuracy rate curve diagram of the application; Figure 6 This is a two-dimensional diagram of a real bearing failure sample from the present invention; Figure 7 This is a two-dimensional diagram of bearing fault samples generated according to the present invention; Figure 8 This is a comparison chart of one-dimensional data of real samples and pseudo samples in this invention; Figure 9 This is a comparison chart of the diagnostic accuracy of different diagnostic algorithms under different sample imbalance rates. Detailed Implementation

[0015] The preferred embodiments of the present invention will now be described in detail with reference to the accompanying drawings, so that the advantages and features of the present invention can be more easily understood by those skilled in the art, thereby providing a clearer and more explicit definition of the scope of protection of the present invention.

[0016] like Figure 1 The data imbalance fault diagnosis method based on improved GAN shown includes the following steps: S1 Data Preparation and Model Improvement: Collect the original imbalanced dataset, generate an adversarial network, and obtain a balanced data generation model that can generate high-quality pseudo-samples by alternating training of the generator and discriminator, thereby improving the adversarial network; S2 Fault Diagnosis Model Construction: High-quality pseudo-samples that closely approximate real samples are generated using an improved adversarial network, and then the pseudo-samples and real samples are input together into a convolutional neural network for fault diagnosis. S3 Experimental Simulation Verification: The generated pseudo-samples are gradually added to the real rolling bearing data to construct a balanced dataset for subsequent fault diagnosis.

[0017] The following is a detailed explanation.

[0018] S1 Data Preparation and Model Improvement S1.1 Constructing the basic generative adversarial network Build a basic GAN model consisting of a generator and a discriminator. The generator is responsible for mapping random noise to generated samples, while the discriminator is responsible for distinguishing whether the input sample comes from real data or the generator, and outputting the probability that it is judged as a real sample.

[0019] Adversarial networks contain two models, one of which is a generative network. The other is to discriminate the network. , Figure 2 The model structure for adversarial networks. Through a generator. and discriminator Through continuous game interaction, the two networks reach a dynamic Nash equilibrium. (Random noise) Sampling from prior distribution , Sampling from real data distribution , The discriminator The activation function of the discriminator The real sample and the generated sample are input into the discriminator , and the output is the probability of judging the source as the real sample. The discriminator is trained to maximize the discriminative loss, while the generator is trained to minimize the loss. The objective function is defined as (1) The discriminative network needs to correctly identify the source of the input data and the generated pseudo data, that is, the probability should be as large as possible. The purpose of the generator is to generate data as close to the real data as possible, that is, the probability should be as small as possible. The training process uses the alternating optimization method of the generator and the discriminator to counter each other. When the output of the discriminator is 0.5, the optimal solution is reached.

[0020] S1.2 Introduce Wasserstein distance and gradient penalty To solve the problems of unstable training, gradient vanishing and mode collapse in standard GAN, the Wasserstein distance is introduced to replace the traditional JS divergence as the loss term for measuring the difference between data distributions. To meet the Lipschitz continuity condition required for the calculation of the Wasserstein distance, a gradient penalty term is introduced in the loss function of the discriminator.

[0021] On the basis of the original generative adversarial network structure, the Wasserstein distance between the input real data and the generated pseudo data distribution is improved as the loss term, and the calculation formula is shown in equation (2). The network parameters of the model are limited within a fixed range to meet the Lipschitz condition, improve the problems of gradient diffusion and mode collapse in the model training process, and ensure the class exclusivity of the generated data. The Lipschitz condition specifies that there is a continuous function that guarantees that the constant K is greater than zero, and any two values , in the entire function domain satisfy equation (3): (2) (3) where represents the probability distribution of the real input data of the real sample, a probability distribution representing generated pseudo data of the generated sample, denotes a mathematical expectation, denotes a distribution function of real input data or generated pseudo data; Because the network parameters of the model are limited in a fixed range, many gradients are distributed on the boundary line, resulting in the problem of gradient disappearance. Gradient penalty is selected to replace the limitation of network parameters in a fixed range, so as to realize the Lipschitz condition and avoid the problem of gradient disappearance caused by improper setting of parameter range. The gradient penalty formula is (4) (5) wherein denotes a cost function of gradient penalty, denotes a penalty coefficient, denotes a linear interpolation calculation on the data of the mixed sample set, as shown in formula (5), , denotes a uniform distribution, denotes a probability distribution between real data and generated data, denotes derivation of , is the output of the network model discriminator, is the output of the network model discriminator, is the 2-norm.

[0022] S1.3 Improved GAN model Through the above improvement, a robust generative adversarial network with Wasserstein distance as the loss and supplemented by gradient penalty mechanism is constructed for subsequent generation of high-quality pseudo samples.

[0023] The generator has two inputs, namely the category label c and the random noise , and the output is the bearing pseudo data. The discriminator has two outputs, in addition to the probability judgment of the data true or false category of the original GAN network, there is also the probability judgment of the label category of the input data. The cost function of the original ACGAN also corresponds to the output structure, first is the loss of the model input data true or false judgment L s , and second is the cross-entropy loss of the input data predicted label and the real label L c . On the basis of the original generative adversarial network structure, the improved one uses the Wasserstein distance of the input real data and the generated pseudo data distribution as the loss term, and the calculation formula is shown in formula (6). Replace the original calculation of the binary cross-entropy function part in formula (3), the model loss function L s and Lc may be expressed as (6) (7) wherein denotes the probability distribution of real bearing data, abbreviated as , denotes the probability distribution of generated pseudo data, abbreviated as , denotes the probability distribution of real label data, denotes the probability distribution of predicted label data, denotes the mathematical expectation, denotes the distribution function of real input data or generated pseudo data.

[0024] The cost function of the improved GAN model is a cost function in which a gradient penalty is added to the cost function of the original ACGAN, and the improved cost function is (8) (9) The discriminant network part needs to output correct labels for the state categories of the generated data or the real data while discriminating the source of the output data, so the loss of the discriminant network part requires maximization , and the generated pseudo data generated by the generator can deceive the discriminant network to a certain extent, leading to misclassification, so the loss of the generator requires maximization .

[0025] S2. Fault diagnosis model construction As shown in Figure 3 , the fault diagnosis model generates high-quality pseudo samples that approximate real samples by using the improved generative adversarial network, and then inputs the pseudo samples and real samples into the convolutional neural network for fault diagnosis to solve the class imbalance problem of bearing data.

[0026] The traditional deep learning model can directly start training after the network is constructed. The generative adversarial network includes two parts, and the model parameters need to be cross-iteratively optimized during training. When training the generative network, the parameters of the discriminant network part also need to be frozen. The specific training process is summarized as four steps: S2.1 Randomly sample random noise in Gaussian distribution , input the random noise and the class label c into the generator to obtain the generated data; S2.2 Input the generated data and the real bearing data into the discriminant network in batches, and the discriminant network calculates the model loss based on the real label, and optimizes the parameters through back propagation; S2.3After each completion of the discriminant network training, the generation network can be trained multiple times. In the early stage of the model, the generation network has weak ability, and the discriminant network can easily distinguish the data source. When the generation network and the discriminant network are combined in series, in order to ensure that the parameter optimization direction of the generation network is correct, the parameters of the discriminant network need to be frozen so that they are not updated with iterations. Then, gradient calculation is performed based on the cost function, and the parameters of the generation network are optimized along the gradient direction; S2.4Repeat steps S2.1-S2.3 to dynamically adjust the training rounds of the generation network according to the loss value fluctuations. When the model reaches Nash dynamic equilibrium, i.e., the pseudo samples generated by the generator can deceive the discriminator, and the discriminator cannot accurately determine the data source but can still correctly classify the bearing fault categories, the model training is completed.

[0027] The model training uses the Adam optimizer. After multiple experiments, the learning rates of the generation network and the discriminant network are set to 0.0002 and 0.0001, respectively. Figure 4 and Figure 5 are the trend curves of the model loss function and accuracy, respectively. From the curves, it can be seen that the model can better distinguish different fault types of bearings at the 2000th epoch, and the classification accuracy is close to 1. The loss curves of the generation network and the discriminant network all reach the set threshold, and the discrimination accuracy of the generated data and the real bearing data fluctuates around 0.5, indicating that the generated samples and the real samples have high similarity, and the discriminator cannot distinguish their sources. The model reaches Nash dynamic equilibrium, and the training is successful.

[0028] After the training is completed, the model is used to generate fault pseudo samples, Figure 6 and Figure 7 respectively show the two-dimensional frequency domain feature maps output by the generator under the real bearing fault categories and the corresponding labels, Figure 8 is a data comparison chart. By converting the two-dimensional frequency domain data of both into one-dimensional sequences and comparing them, it can be seen that the generated samples and the real samples have high similarity.

[0029] The analysis results show that the pseudo samples output by the generation network are basically consistent with the real samples in the main features, and the key data distribution areas of each fault category are prominent. The generated data is highly consistent with the real samples in amplitude and time sequence trend. However, in the area where the feature of the real sample is not obvious (such as the area where the vertical coordinate value is abnormally high or low), the generated sample does not completely reproduce its distribution characteristics. This difference effectively alleviates the problem of data repetition commonly seen in oversampling, avoids the risk of overfitting caused by the excessive similarity of expanded samples, and reflects the good robustness of the proposed fault diagnosis model. In summary, the model training is basically successful, can fully excavate the deep features of the fault data, and ensures the mutual exclusivity of the generated samples, laying a foundation for solving the class imbalance problem through data expansion in the future.

[0030] S3 Experimental Simulation Verification The trained model is then applied to generate high-quality pseudo-samples. To simulate real-world industrial class imbalance scenarios, these pseudo-samples are gradually added to real rolling bearing data to construct a balanced dataset for subsequent fault diagnosis. In the simulation experiments, it is assumed that the majority class (…) That is, the number of bearing samples under normal operating conditions is 1000, unbalanced type ( This means the number of bearing samples that failed under fault conditions is 10. The sample imbalance rate is defined. The formula is (10) Specify the initial sample imbalance rate The ratio is 1:100. Fault diagnosis tests were conducted at ratios of 1:50, 1:25, 1:10, and 1:1, respectively. Specifically, for each type of bearing fault sample, data augmentation was performed using generated pseudo-samples (the augmented training set is shown in Table 1). Based on this, multiple bearing fault diagnosis models were trained. Figure 9 The diagnostic accuracy of different diagnostic algorithms under different sample imbalance rates is shown for performance comparison. Figure 9 The results show that as pseudo-samples are continuously injected into the originally extremely imbalanced training set, the accuracy of the five mainstream intelligent diagnostic methods for rolling bearings steadily climbs to above 0.9. It should be noted that none of the listed networks underwent fine-tuning for this bearing dataset; only default configurations commonly used in the fault diagnosis field were employed. When the ratio of positive to negative samples became 1:100, the performance of each model was generally poor. SAE almost completely lost its fault identification ability, primarily because the extremely scarce fault samples prevented the model from fully exploring the discriminative features of the minority class, leading to the vast majority of fault samples being misclassified as healthy. However, when the samples were replenished to a 1:1 balance, the recognition accuracy of BPNN, LeNet5, and ResNet18 all approached 0.99.

[0031] The above trends indicate that as the nine types of fault samples are gradually supplemented, the network can learn effective features from the expanded data; even the introduction of spurious samples did not cause training collapse, indicating that the distribution of the generated samples is highly consistent with the real data. The subsequent five diagnostic models can all perform deep representation learning for various faults, and with the increase in data scale and balance, the classification accuracy also increases synchronously.

[0032] The above are all preferred embodiments of the present application, and do not limit the protection scope of the present application, so that: all equivalent changes made according to the structure, shape, principle and application direction of the present application should be covered in the protection scope of the present application.

Claims

1. A data imbalance fault diagnosis method based on an improved GAN, characterized in that, Includes the following steps: S1 Data Preparation and Model Improvement: Collect the original imbalanced dataset, generate an adversarial network, and obtain a balanced data generation model that can generate high-quality pseudo-samples by alternating training of the generator and discriminator, thereby improving the adversarial network; S2 Fault Diagnosis Model Construction: High-quality pseudo-samples that closely approximate real samples are generated using an improved adversarial network, and then the pseudo-samples and real samples are input together into a convolutional neural network for fault diagnosis. S3 Experimental Simulation Verification: The generated pseudo-samples are gradually added to the real rolling bearing data to construct a balanced dataset for subsequent fault diagnosis.

2. The data imbalance fault diagnosis method based on improved GAN according to claim 1, characterized in that, Step S1 specifically includes the following steps: S1.1 Constructing a basic generative adversarial network: Build a basic GAN model consisting of a generator and a discriminator; where the generator is responsible for mapping random noise to generated samples; the discriminator is responsible for distinguishing whether the input sample comes from real data or the generator, and outputs the probability of it being judged as a real sample; S1.2 Introduction of Wasserstein distance and gradient penalty: Wasserstein distance is introduced as a loss term to measure the difference in data distribution. To meet the Lipschitz continuity condition required for Wasserstein distance calculation, a gradient penalty term is introduced into the loss function of the discriminator. S1.3 Improved GAN model: Construct a robust generative adversarial network with Wasserstein distance as loss and gradient penalty mechanism for the generation of high-quality pseudo samples.

3. The data imbalance fault diagnosis method based on improved GAN according to claim 2, characterized in that, In step S1.1, random noise Sampling from prior distribution Real samples Sampling from real data distribution , To generate samples, train the discriminator. Maximize the discriminative loss while training the generator. To minimize the loss, the objective function is defined as follows: （1）。 4. The data imbalance fault diagnosis method based on improved GAN according to claim 3, characterized in that, In step S1.2, the formula for calculating the loss term using the Wasserstein distance is as follows: （2） The Lipschitz condition specifies the existence of a continuous function that guarantees a constant K greater than zero, and that any two values ​​of K within the entire domain of the function are valid. , satisfy: （3） in This represents the probability distribution of the real input data for the real samples. This represents the probability distribution of generated pseudo-data for the generated samples. Represents the mathematical expectation. A distribution function representing real input data or generated pseudo data; Gradient penalty is used instead of restricting network parameters to a fixed range to achieve the Lipschitz condition and avoid gradient vanishing caused by improper parameter range settings. The gradient penalty formula is as follows: （4） （5） in The cost function represents the gradient penalty. Indicates the penalty coefficient. This indicates that the data from the mixed sample set is obtained through linear interpolation. , Indicates uniform distribution. This represents the probability distribution between real input data and generated pseudo data. Indicates to Differentiate, yes The output after passing through the network model discriminator It is a 2-norm.

5. The data imbalance fault diagnosis method based on improved GAN according to claim 4, characterized in that, In step S1.3, the Wasserstein distance between the distributions of the input real data and the generated pseudo-data is improved and used as the loss term, replacing the original binary cross-entropy function calculation part. The model loss function... L s and L c They can be represented as follows: （6） （7） in This represents the probability distribution of the actual input data. This represents the probability distribution for generating pseudo-data. This represents the probability distribution of the real labeled data. This represents the probability distribution of the predicted label data; The improved cost function of an adversarial network is obtained by adding a gradient penalty cost function to the original cost function. The improved cost function is then: （8） （9） 。 6. The data imbalance fault diagnosis method based on improved GAN according to claim 1, characterized in that, Step S2 specifically includes the following steps: S2.1 Random noise is generated by randomly sampling from a Gaussian distribution. random noise Input the category label 'c' into the generator to obtain the generated data; S2.2 The generated data and real bearing data are input into the discrimination network in batches. The discrimination network calculates the model loss based on the real labels and optimizes the parameters through backpropagation. S2.3 After each training iteration of the discriminator network is completed, the generator network can be trained multiple times. The parameters of the discriminator network are frozen so that they are not updated with iteration. Then, gradient calculation is performed based on the cost function, and the parameters of the generator network are optimized along the gradient direction. S2.4 Repeat steps S2.1-S2.3, dynamically adjusting the number of training rounds of the generator network based on the fluctuation of the loss value; when the model reaches Nash dynamic equilibrium, that is, when the pseudo-samples generated by the generator can deceive the discriminator, and the discriminator cannot accurately determine the source of the data but can still correctly classify the bearing fault category, the model training ends.

Images