Gearbox fault diagnosis method based on spatial transformation network and attention mechanism

Abstract

The present application relates to a kind of gearbox fault diagnosis methods based on spatial transformation network and attention mechanism, it includes the following steps, step 1: the vibration signal data of gear box is collected, one-dimensional vibration signal data is converted into two-dimensional time-frequency picture data;Step 2: the two-dimensional time-frequency picture data group of gear box is grouped, obtains training set, verification set and test set;Step 3: construct gearbox fault diagnosis model;Step 4: complete the training of gearbox fault diagnosis model, output diagnosis result.The present application adds spatial transformation network and attention mechanism in the process of constructing gearbox intelligent fault diagnosis model, thereby improve the geometric invariability of model, allow network to selectively enhance the feature of a large amount of information, so that subsequent processing can make full use of these features, so that model is easier to extract the feature of fault, realizes intelligent fault diagnosis;Using multi-label class balance loss function, the problem of unbalanced data training is solved.

Description

Technical Field

[0001] The present invention relates to the field of mechanical equipment fault detection and intelligent diagnosis, and in particular to a gearbox fault diagnosis method based on a spatial transformation network and an attention mechanism. Background Art

[0002] Gearboxes are widely used in a variety of modern industrial machinery, including wind turbines, automobiles, and aircraft, and are crucial transmission components in many modern machines. Because gearboxes often operate in harsh environments such as high temperature and humidity, they are susceptible to a wide range of failures. Failure to promptly diagnose and identify the fault type can have catastrophic consequences, making gearbox fault diagnosis essential.

[0003] With the continuous development of machine learning, machine learning techniques are increasingly being applied in the field of fault diagnosis. For example, intelligent fault diagnosis technology, which combines convolutional neural networks with fault diagnosis, has been widely used. This method uses methods such as Fourier transform to extract features from raw vibration signals and then uses convolutional neural networks for fault classification, achieving good accuracy. However, for actual application in production, two issues remain to be addressed.

[0004] First, there's the issue of data imbalance. In actual production, mechanical equipment operates normally most of the time, with only very short periods of failure. Furthermore, due to varying operating environments and conditions, the types of mechanical failures vary, and the probabilities of various failure types vary. Therefore, in actual production, the number of samples representing different health states varies. While some methods achieve good accuracy using standard datasets, they may not necessarily translate well in real-world production. Second, there's the multi-label issue. In actual production, multiple faults can occur simultaneously. Using a single label would effectively treat these combined faults as a single, nonsensical type. Therefore, given the data imbalance, leveraging deep learning to achieve high-precision diagnosis of combined gearbox faults is a pressing challenge. Summary of the Invention

[0005] In order to overcome the shortcomings of the existing technology, the present invention adds a spatial transformation network and an attention mechanism in the process of constructing an intelligent fault diagnosis model for gearboxes, thereby improving the geometric invariance of the model and allowing the network to selectively enhance the features of a large amount of information, so that subsequent processing can fully utilize these features, making it easier for the model to extract fault features and realize intelligent fault diagnosis; using a multi-label class balanced loss function, the problem of unbalanced data training is solved.

[0006] To achieve the above objectives, the solution adopted by the present invention is:

[0007] A gearbox fault diagnosis method based on a spatial transformer network and an attention mechanism comprises the following steps:

[0008] Step 1: Collect the vibration signal data of the gearbox and convert the one-dimensional vibration signal data into two-dimensional time-frequency image data;

[0009] The collected vibration signal is divided into continuous data segments, and the wavelet transform is used to transform the one-dimensional vibration signal data into two-dimensional time-frequency image data, as shown below:

[0010]

[0011] Where: d represents the two-dimensional time-frequency image data set generated after wavelet transformation, specifically d = [d1, d2, ...d s ]; z(t) represents the vibration signal collected at time t; represents the wavelet basis function; a represents the scale factor; τ represents the translation amount; t is the time of collecting vibration signals; s represents the total number of images; d1, d2, ... d s Respectively represent the 1st, 2nd…sth two-dimensional time-frequency image data groups;

[0012] Step 2: Group the two-dimensional time-frequency image data set of the gearbox to obtain a training set, a validation set, and a test set;

[0013] Obtain the two-dimensional time-frequency image data set from step 1, and divide the two-dimensional time-frequency image data set into a ratio of 7:2:1 to obtain a training set, a validation set, and a test set, respectively; delete the two-dimensional time-frequency image data of the gearbox working state in the training set to obtain multiple training sets with unbalanced data;

[0014] Step 3: Construct a gearbox fault diagnosis model;

[0015] Step 31: Construct a spatial transformation network, which includes a positioning network, a grid generator, and a sampler. Input the two-dimensional time-frequency image data into the spatial transformation network for processing to improve geometric invariance. The output result is as follows:

[0016] layer1 = STN(d);

[0017] Where: layer1 represents the output of the spatial transformer network; STN represents the spatial transformer network;

[0018] Step 32: Construct an improved deep residual network model, and input the calculation results of step 31 into the improved deep residual network. The prediction results of the improved deep residual network are as follows:

[0019] m=model(layer1);

[0020] Where: m=[m1,m2,…,m C ] represents the prediction output of the improved deep residual network model; m1, m2, ..., m C They represent the prediction output results of the first, second, ...Cth improved deep residual network models respectively; model represents the improved deep residual network;

[0021] Step 33: Add a loss function and an activation function to the improved deep residual network model; the output of the loss function is as follows:

[0022]

[0023] Where: CB sigmoid (m,y) represents the output of the loss function; y represents the actual label of the category; β represents a hyperparameter; n y Indicates the number of training samples with label y; m i Represents the predicted output result of the i-th improved deep residual network model; i represents the predicted output result number of the improved deep residual network model, i∈[1,2,…,c];

[0024] The predicted output result m of the i-th improved deep residual network model i As shown below:

[0025]

[0026] The activation function expression is as follows:

[0027] f(u) = u·sigmoid(u);

[0028] Where: f(u) represents the output of the activation function; sigmoid represents the activation function; u represents the input of the activation function;

[0029] Step 4: Complete the training of the gearbox fault diagnosis model and output the diagnosis results;

[0030] Repeat step 3 to train the unbalanced training set in step 2, and use the validation set for verification, saving the gearbox fault diagnosis model with the highest accuracy; test it based on the test set and output the final fault diagnosis results.

[0031] Preferably, in step 2, the two-dimensional time-frequency image data of the gearbox working state in the training set is deleted to obtain multiple training sets with unbalanced data, specifically:

[0032] The gearbox working state includes a normal working state, an inner ring fault state and an outer ring fault state;

[0033] The number of samples of each class in the unbalanced dataset follows a geometric distribution, and the ratio of the maximum number of samples to the minimum number of samples is im, where im represents the class imbalance factor. Four sets of datasets with different im values ​​are set, and the values ​​of im are 1, 35, 70, and 140, respectively.

[0034] Preferably, the positioning network in step 31 is specifically:

[0035] The positioning network is a regression network that converts the time-frequency graph into a regression parameter θ through a hidden layer. The expression is as follows:

[0036] θ=f loc (U); U∈R G×H×W ;

[0037] Where: θ represents the output vector of the positioning network; f loc Represents the positioning network function; U represents the time-frequency diagram of the positioning network input; R represents a natural number; G, H and W represent the number of channels, length and width of the positioning network respectively.

[0038] Preferably, the grid generator in step 31 is specifically:

[0039] The grid generator is a sampling grid constructed according to the transformation parameters, which can obtain the coordinate points of the output feature map corresponding to the coordinate points of the input feature map; given the source coordinates in the input feature map And the learned two-dimensional affine transformation matrix A θ , output the target coordinates of the regular grid in the feature map As shown below;

[0040]

[0041] Where: and Represent the horizontal and vertical coordinates in the input feature map respectively; and Respectively represent the target horizontal coordinate and vertical coordinate of the regular grid in the output feature map; A θ represents the two-dimensional affine transformation matrix; θ 11 ,θ 12 …θ 23 They represent the 11th, 12th, ..., 23rd elements of the two-dimensional affine transformation matrix respectively.

[0042] Preferably, the sampler in step 31 is specifically:

[0043] The sampler performs bilinear sampling on the input time-frequency map U and the sampling grid to generate a converted output feature map V. Through the spatial transformation layer, the image data can be sheared, translated, scaled and rotated to improve the accuracy of classification.

[0044] Preferably, the improved deep residual network model in step 32 is specifically:

[0045] The improved deep residual network includes a convolutional layer, a batch normalization layer, a maximum pooling layer, an attention mechanism module and a global average pooling layer;

[0046] The attention mechanism allows the network to selectively enhance the features of a large amount of information, enabling subsequent processing to fully utilize these features and suppress useless features. This assigns different weights to different channels, making it easier for the model to extract fault features and thus achieve intelligent fault diagnosis. The attention module consists of 10 stages, each with the same structure but different model parameters.

[0047] The global average pooling layer can suppress overfitting and make the input size more flexible.

[0048] Compared with the prior art, the present invention has the following beneficial effects:

[0049] (1) In the process of constructing the gearbox intelligent fault diagnosis model, the present invention adds a spatial transformation network, thereby improving the geometric invariance of the model and enabling the model to achieve the functions of changing specific areas, amplifying images, and correcting image directions;

[0050] (2) In the process of constructing the gearbox intelligent fault diagnosis model, the present invention incorporates an attention mechanism, which allows the network to selectively enhance the features of a large amount of information, so that subsequent processing can fully utilize these features and suppress useless features, making it easier for the model to extract fault features, thereby realizing intelligent fault diagnosis;

[0051] (3) The classification method adopted by the present invention is a multi-label classification method, which can diagnose the composite faults of the gearbox. Compared with the single-label classification method, the multi-label classification is more accurate and more in line with the actual production process;

[0052] (4) The present invention uses a multi-label class balanced loss function to calculate the effective sample number by associating each sample to a small adjacent area rather than a single point; the proposed weighting factor is inversely proportional to the effective sample capacity, and the problem of unbalanced data training is solved by introducing the weighting factor. BRIEF DESCRIPTION OF THE DRAWINGS

[0053] Figure 1 This is a flowchart of a gearbox fault diagnosis method based on a spatial transformation network and an attention mechanism according to an embodiment of the present invention;

[0054] Figure 2 This is a flow chart of a fault diagnosis method for a gearbox according to an embodiment of the present invention;

[0055] Figure 3 It is the experimental platform in the example of the present invention;

[0056] Figure 4 Schematic diagram of the improved deep residual network structure in an example of the present invention;

[0057] Figure 5 This is the spatial transformation network structure of the present invention;

[0058] Figure 6 Locating the network structure for the present invention example;

[0059] Figure 7 This is the example attention mechanism module structure of the present invention. DETAILED DESCRIPTION

[0060] Hereinafter, embodiments of the present invention will be described with reference to the accompanying drawings.

[0061] The embodiment of the present invention is based on wavelet transform and improved deep residual network model, uses wavelet transform for feature extraction, then adds spatial transformation network layer and attention mechanism to the deep residual network, and uses multi-label class balance loss function to realize gearbox composite fault diagnosis under data imbalance, making it easier for the model to extract fault features, thereby realizing intelligent fault diagnosis. Figure 1 Shown is a flowchart of a gearbox fault diagnosis method based on a spatial transformation network and an attention mechanism according to an example of the present invention.

[0062] The embodiment of the present invention provides a gearbox fault diagnosis method based on a spatial transformation network and an attention mechanism, such as Figure 2 The flowchart of the fault diagnosis method of the gearbox of the present invention is shown. In order to demonstrate the applicability of the present invention, the method is applied to an example, which specifically includes the following steps:

[0063] S1: Collect the vibration signal data of the gearbox and convert the one-dimensional vibration signal data into two-dimensional time-frequency image data;

[0064] like Figure 3 The figure shows the experimental platform used for data collection in this patent. 1 is the prime mover, 2 is the gear, 3 is the flywheel, 4 is the asynchronous motor, and 5 is the computer. The experimental platform primarily includes a gearbox, prime mover, voltage regulator, asynchronous generator, and flywheel. Two triaxial accelerometers are mounted on the outside of the gearbox to collect horizontal and vertical vibration signal data.

[0065] Defects on the outer ring, balls, and inner ring of the bearing are all caused by electric sparks. Shims are used to adjust the height of the transmission to simulate mechanical misalignment of the transmission. There are 16 health conditions, including normal conditions, 7 single fault types, and 8 combined fault types. Single faults include: broken teeth, gear shedding, gear cracks, bearing outer ring damage, bearing ball damage, loose gearbox fixed plate, and axial imbalance. Combined faults include: broken teeth & gear shedding, broken teeth & gear cracks, broken teeth & bearing outer ring damage, broken teeth & loose fixed plate, gear shedding & bearing outer ring damage, gear shedding & bearing ball damage, gear shedding & loose fixed plate, loose fixed plate & axial imbalance.

[0066] In this embodiment, the sampling points of each fault mode are divided into 400 data segments with a segment length of 2600 sampling points. That is, each data segment contains 2600 continuous sampling points. The collected vibration signal is divided into continuous data segments, and the one-dimensional vibration signal data is converted into two-dimensional time-frequency image data using wavelet transform, as shown below:

[0067]

[0068] Where: d represents the two-dimensional time-frequency image data set generated after wavelet transformation, specifically d = [d1, d2, ...d s ]; z(t) represents the vibration signal collected at time t; represents the wavelet basis function, which is Morlet wavelet in this embodiment, and the center frequency and bandwidth are both 3; a represents the scale factor; τ represents the translation amount; t is the time for collecting vibration signals; s represents the total number of images; d1, d2, ... d s They represent the first, second, ... sth two-dimensional time-frequency image data groups respectively.

[0069] S2: Group the two-dimensional time-frequency image data set of the gearbox to obtain a training set, a validation set, and a test set;

[0070] A two-dimensional time-frequency image data set of S1 was obtained and divided into a training set, a validation set, and a test set in a ratio of 7:2:1. The two-dimensional time-frequency image data of the gearbox working state in the training set was deleted to obtain multiple training sets with unbalanced data. These sets were all normalized to [-1, 1]. In the standard data set, for each fault mode, the training set includes 280 time-frequency images, which is the number of training samples; the validation set includes 80 time-frequency images, which is the number of validation samples; and the test set includes 40 time-frequency images, which is the number of test samples. Each time-frequency image is a sample. Detailed data is shown in Table 1, which shows the training set and sample set of wavelet time-frequency images for each fault mode.

[0071] Table 1: Training set and sample set of wavelet time-frequency images for each fault mode

[0072]

[0073] The gearbox working state includes normal working state, inner ring fault state and outer ring fault state.

[0074] The number of samples of each class in an unbalanced data set follows a geometric distribution. The ratio of the maximum number of samples to the minimum number of samples is im, where im represents the class imbalance factor. Four sets of data sets with different im values ​​are set, and the values ​​of im are 1, 35, 70, and 140 respectively.

[0075] S3: Construct a gearbox fault diagnosis model;

[0076] S31: Build a spatial transformation network, which includes a positioning network, a grid generator, and a sampler. Input the two-dimensional time-frequency image data into the spatial transformation network for processing to improve geometric invariance. The output results are as follows:

[0077] layer1 = STN(d);

[0078] Where: layer1 represents the output result of the spatial transformer network; STN represents the spatial transformer network.

[0079] The localization network is a regression network that converts the time-frequency graph into the regression parameter θ through the hidden layer. The expression is as follows:

[0080] θ=f loc (U); U∈R G×H×W ;

[0081] Where: θ represents the output vector of the positioning network; f loc Represents the positioning network function; U represents the time-frequency diagram of the positioning network input; R represents a natural number; G, H and W represent the number of channels, length and width of the positioning network respectively.

[0082] The grid generator is a sampling grid constructed according to the transformation parameters, which can obtain the coordinate points of the output feature map corresponding to the coordinate points of the input feature map; given the source coordinates in the input feature map And the learned two-dimensional affine transformation matrix A θ , output the target coordinates of the regular grid in the feature map As shown below:

[0083]

[0084] Where: and Represent the horizontal and vertical coordinates in the input feature map respectively; and Respectively represent the target horizontal coordinate and vertical coordinate of the regular grid in the output feature map; A θ represents the two-dimensional affine transformation matrix; θ 11 ,θ 12 …θ 23 They represent the 11th, 12th, ..., 23rd elements of the two-dimensional affine transformation matrix respectively.

[0085] The sampler performs bilinear sampling on the input time-frequency map U and the sampling grid to generate a converted output feature map V. Through the spatial transformation layer, the image data can be sheared, translated, scaled and rotated to improve the accuracy of classification.

[0086] S32: Build an improved deep residual network model and feed the calculation results of S31 into the improved deep residual network. The prediction results of the improved deep residual network are as follows:

[0087] m=model(layer1);

[0088] Where: m=[m1,m2,…,m C ] represents the prediction output of the improved deep residual network model; m1, m2, ..., m C They represent the predicted output results of the first, second, ...Cth improved deep residual network models respectively; model represents the improved deep residual network.

[0089] The method for creating an improved fault diagnosis model includes: designing a structure of a deep residual network based on a neural network model.

[0090] Specifically, the traditional deep residual network structure is improved, mainly including adding a spatial transformation network layer; using an attention mechanism module to replace the traditional residual module; and using global average pooling to replace the traditional fully connected layer.

[0091] like Figure 4 Figure 2 shows the improved deep residual network architecture used in this example. The deep residual network includes multiple convolutional layers, multiple pooling layers, an attention module, and global average pooling. The attention mechanism allows the network to selectively enhance features that contain a large amount of information, enabling subsequent processing to fully utilize these features and suppress unhelpful ones, thereby assigning different weights to different channels. This paper combines the attention mechanism with a deep residual network, making it easier for the model to extract fault characteristics.

[0092] like Figure 5 The example space transformation network structure of the present invention is shown as follows; it includes Figure 6The localization network structure, grid generator, and sampler are shown. The localization network is a simple regression network. The grid generator constructs a sampling grid based on the transformation parameters, which can obtain the coordinate points of the output feature map corresponding to the coordinate points of the input feature map. The sampler performs bilinear sampling on the input feature map U and the sampling grid to produce the transformed output feature map V.

[0093] like Figure 7 The following figure shows the structure of the attention mechanism module in an example of this invention. Global average pooling is first used to compress the spatial dimensions to 1×1. A convolutional layer is then used to compress the C channel into a C / R channel. Silu activation is then applied, followed by a convolutional layer to restore the number of channels to C. Finally, a sigmoid activation function is used to multiply the result by the original channel to achieve weight distribution.

[0094] S33: Add a loss function and an activation function to the improved deep residual network model; the embodiment of the present invention uses a multi-label class balanced logistic loss function. Compared with the traditional logistic loss function, the improved loss function has stronger robustness; the output of the loss function is as follows:

[0095]

[0096] Where: CB sigmoid (m,y) represents the output of the loss function; y represents the actual label of the category; β represents a hyperparameter; n y Indicates the number of training samples with label y; m i Represents the predicted output result of the i-th improved deep residual network model; i represents the predicted output result number of the improved deep residual network model, i∈[1,2,…,c].

[0097] The predicted output result m of the i-th improved deep residual network model i As shown below:

[0098]

[0099] The activation function expression is as follows:

[0100] f(u) = u·sigmoid(u);

[0101] Where: f(u) represents the output of the activation function; sigmoid represents the activation function; u represents the input of the activation function.

[0102] S4: Complete the training of the gearbox fault diagnosis model and output the diagnosis results;

[0103] Repeat S3 to train the unbalanced training set in S2, and use the validation set for verification, saving the gearbox fault diagnosis model with the highest accuracy; test according to the test set, and output the final fault diagnosis result.

[0104] In order to verify the generalization and robustness of the present invention, five repeated experiments were conducted, and the experimental results are shown in Table 2.

[0105] Table 2 Results of 5 experiments and average test accuracy

[0106]

[0107] In actual production processes, the number of different health states is unequal; the number of healthy states far exceeds that of other types. Furthermore, due to varying working environments and conditions, the probability of various machine failures also varies. Here, we created three unbalanced datasets. Four sets of experiments were conducted, each repeated five times to minimize errors caused by experimental contingency. The experimental results are shown in Table 2. The results of the 25 experiments and the average test accuracy in Table 2 show that when im = 1 (in the standard dataset), although the diagnostic accuracy of the traditional deep residual network model (WT-TDRN) is lower than that of the proposed model (WT-IDRN), both achieve diagnostic accuracy exceeding 99%. However, as the imbalance factor increases, the average accuracy of both approaches decreases. When im = 140, the diagnostic accuracy of the WT-TDRN lags significantly behind that of the proposed WT-IDRN. This demonstrates that the proposed method can effectively address data imbalance and complex fault diagnosis issues.

[0108] Table 3 Diagnostic accuracy of each health status in the experiment

[0109]

[0110] Table 3 shows the diagnostic accuracy of each health state. While the WT-TDRN model has good overall diagnostic accuracy, it has particularly low accuracy for certain specific fault categories. If, in actual production, the probability of a gearbox experiencing specific faults is high due to the production environment and operating conditions, and the WT-TDRN model has a low diagnostic accuracy for these faults, catastrophic consequences could result.

[0111] To more clearly demonstrate the advantages of the present invention's multi-label approach to complex fault diagnosis, we compared the results of the fault diagnosis method of an embodiment of the present invention with those of other advanced methods using a standard dataset, based on gearbox test data from a university. The results are shown in Table 4.

[0112] Table 4 Comparison of the average accuracy of the present invention (WT-IDRN) with other methods

[0113]

[0114] From the comparison of the average accuracy of the present invention (WT-IDRN) and other methods in Table 4, it can be seen that compared with the existing methods, the method proposed in this paper is better and has the highest diagnostic accuracy, reaching 99.38%. This shows that the model proposed in this paper is very effective in composite fault diagnosis.

[0115] In summary, the prediction results of the gearbox fault diagnosis method based on spatial transformer network and attention mechanism in this case have proved to be very effective.

[0116] (1) In the process of constructing the intelligent fault diagnosis model of gearboxes, the embodiment of the present invention adds a spatial transformation network and an attention mechanism; improves the geometric invariance of the model, so that the model can change specific areas, amplify images and correct image directions; can selectively enhance the features of a large amount of information, so that subsequent processing can make full use of these features, suppress useless features, and make it easier for the model to extract fault features, thereby realizing intelligent fault diagnosis; through the computational analysis of the embodiment, it can be seen that this case has significantly improved the accuracy of gearbox diagnosis, and the actual application effect is good.

[0117] (2) The classification method adopted in the embodiment of the present invention is a multi-label classification method, which can diagnose the complex faults of the gearbox. Compared with the single-label classification method, the multi-label classification is more accurate and more in line with the actual production and manufacturing process. The multi-label class balance loss function is used to calculate the number of effective samples by associating each sample to a small adjacent area rather than a single point. The weighting factor proposed in this case is inversely proportional to the effective sample capacity. The problem of unbalanced data training is solved by introducing the weighting factor.

[0118] (3) Compared with traditional methods, the method proposed in the embodiment of the present invention is more practical and has higher practical application value; it provides a new, feasible and effective way to solve the problem of gearbox fault diagnosis, and can be widely used in fault diagnosis in the fields of machinery, chemical industry, aerospace, etc.

[0119] The embodiments described above are merely descriptions of preferred implementations of the present invention and are not intended to limit the scope of the present invention. Without departing from the spirit of the present invention, various modifications and improvements made to the technical solutions of the present invention by ordinary technicians in this field should fall within the scope of protection determined by the claims of the present invention.

Claims

1. A gearbox fault diagnosis method based on spatial transformation network and attention mechanism, characterized in that: It includes the following steps: Step 1: Collect the vibration signal data of the gearbox and convert the one-dimensional vibration signal data into two-dimensional time-frequency image data; The collected vibration signal is divided into continuous data segments, and the wavelet transform is used to transform the one-dimensional vibration signal data into two-dimensional time-frequency image data, as shown below: ； Where: Represents a two-dimensional time-frequency image data set generated after wavelet transform, specifically ; represents the vibration signal collected at time t; represents the wavelet basis function; represents the scale factor; Indicates the amount of translation; is the time of collecting vibration signals; s represents the total number of pictures; Respectively represent the 1st, 2nd…sth two-dimensional time-frequency image data groups; Step 2: Group the two-dimensional time-frequency image data set of the gearbox to obtain a training set, a validation set, and a test set; Obtain the two-dimensional time-frequency image data set from step 1, and divide the two-dimensional time-frequency image data set into a ratio of 7:2:1 to obtain a training set, a validation set, and a test set, respectively; delete the two-dimensional time-frequency image data of the gearbox working state in the training set to obtain multiple training sets with unbalanced data; Step 3: Construct a gearbox fault diagnosis model; Step 31: Construct a spatial transformation network, which includes a positioning network, a grid generator, and a sampler. Input the two-dimensional time-frequency image data into the spatial transformation network for processing. The output result is as follows: ； Where: Represents the output of the spatial transformation network; Represents a spatial transformation network; The positioning network is a regression network that converts the time-frequency graph into regression parameters through the hidden layer. , the expression is as follows: ； Where: Represents the output vector of the localization network; Represents the positioning network function; U represents the time-frequency diagram of the positioning network input; R represents a natural number; G, H and W represent the number of channels, length and width of the positioning network respectively; Step 32: Construct an improved deep residual network model, and input the calculation results of step 31 into the improved deep residual network. The prediction results of the improved deep residual network are as follows: ； Where: represents the prediction output of the improved deep residual network model; Represent the prediction output results of the 1st, 2nd...Cth improved deep residual network models respectively; Represents the improved deep residual network; The improved deep residual network includes a convolutional layer, a batch normalization layer, a maximum pooling layer, an attention mechanism module and a global average pooling layer; The attention mechanism module consists of 10 stages, each with the same structure but different model parameters; Step 33: Add a loss function and an activation function to the improved deep residual network model; the output of the loss function is as follows: ； Where: Represents the output of the loss function; The actual label representing the category; represents a hyperparameter; Indicates that the label is The number of training samples; m i represents the predicted output result of the i-th improved deep residual network model; i represents the predicted output result number of the improved deep residual network model, ; The predicted output result m of the i-th improved deep residual network model i As shown below: ； The activation function expression is as follows: ； Where: Represents the output of the activation function; represents the activation function; Represents the input of the activation function; Step 4: Complete the training of the gearbox fault diagnosis model and output the diagnosis results; Repeat step 3 to train the unbalanced training set in step 2 and use the validation set for verification; test the system based on the test set and output the final fault diagnosis results.

2. The gearbox fault diagnosis method based on spatial transformation network and attention mechanism according to claim 1 is characterized in that: In step 2, the two-dimensional time-frequency image data of the gearbox working state in the training set is deleted to obtain multiple training sets with unbalanced data, specifically: The gearbox working state includes a normal working state, an inner ring fault state and an outer ring fault state; The number of samples of each class in the unbalanced dataset follows a geometric distribution, the ratio of the maximum number of samples to the minimum number of samples is im, and im represents the class imbalance factor. Four sets of datasets with different im values ​​are set, and the values ​​of im are 1, 35, 70, and 140 respectively.

3. The gearbox fault diagnosis method based on spatial transformation network and attention mechanism according to claim 1 is characterized in that: The grid generator in step 31 is specifically: The grid generator is a sampling grid constructed according to the transformation parameters, which can obtain the coordinate points of the output feature map corresponding to the coordinate points of the input feature map; given the source coordinates in the input feature map And the learned two-dimensional affine transformation matrix , output the target coordinates of the regular grid in the feature map , as shown below; ； Where: Represent the horizontal and vertical coordinates in the input feature map respectively; They represent the target horizontal and vertical coordinates of the regular grid in the output feature map respectively; Represents a two-dimensional affine transformation matrix; 、 … They represent the 11th, 12th, ..., 23rd elements of the two-dimensional affine transformation matrix respectively.

4. The gearbox fault diagnosis method based on spatial transformation network and attention mechanism according to claim 1 is characterized in that: The sampler in step 31 is specifically: The sampler performs bilinear sampling on the input time-frequency map U and the sampling grid to generate a converted output feature map V. Through the spatial transformation layer, the image data can be sheared, translated, scaled and rotated.

Images