Gear fault recognition method based on improved multi-head attention

Abstract

The application discloses a gear fault identification method based on an improved multi-head attention mechanism, and relates to fault identification. Vibration signals of a gear in different fault states are acquired. The vibration signals are preprocessed to obtain a standardized sample set. A double-path convolution network is used to extract and fuse features of samples in the standardized sample set, so as to obtain a fused feature sequence. An improved multi-head attention mechanism is used to enhance features of the fused feature sequence, so as to obtain an attention feature sequence. The attention feature sequence is subjected to time sequence feature coding and global feature aggregation, so as to obtain global features. Gear fault categories are identified according to the global features. In view of the fact that the multi-head attention mechanism in the prior art has high calculation complexity when processing one-dimensional time sequence data of vibration signals, the application effectively reduces the calculation complexity.

Description

Technical Field

[0001] This application relates to the field of fault diagnosis, and more specifically, to a gear fault identification method based on improved multi-head attention. Background Technology

[0002] With the continuous upgrading and intelligent transformation of the manufacturing industry, higher requirements are being placed on the condition monitoring and health management of mechanical equipment. As the core of complex transmission systems such as planetary gearboxes, the operating status of gears directly determines the reliability and safety of the entire equipment. In actual production, gears are subjected to complex and variable working conditions such as high speed and heavy load for extended periods, making them prone to various typical faults such as pitting, cracking, and tooth breakage. These can lead to equipment downtime or even major production accidents. Therefore, conducting accurate and efficient gear health condition monitoring and fault diagnosis analysis is extremely important.

[0003] Traditional fault diagnosis relies heavily on signal processing and machine learning, which often struggle to extract deep information from fault features, resulting in limited recognition accuracy. Therefore, deep learning-based intelligent diagnostic methods are frequently employed to improve the accuracy of fault identification and classification. However, the increased complexity of the model structure and the expansion of the number of parameters also bring corresponding problems. Especially when processing long input sequences, the model generates a massive amount of intermediate activation values ​​and cached data, causing a sharp increase in GPU memory usage, leading to training delays or even recognition failures, severely impacting the real-time performance and deployment feasibility of the diagnostic system.

[0004] To overcome the aforementioned bottlenecks, this application proposes a multi-module caching optimization method. Dual-path convolution is used to extract global features of the signal, extending common one-dimensional fault signals to multi-dimensional ones to facilitate head number partitioning in multi-head latent attention. Then, by improving the multi-head latent attention mechanism, a low-rank projection sharing mechanism is used to reuse parameters in the attention module, reducing cache usage during computation. Simultaneously, a sine / cosine position encoding compatibility design is employed to increase the model's adaptability to fault signals, avoiding parameter sharing compatibility issues caused by the introduction of position information. Finally, a bidirectional gated recurrent unit is used to fully extract the temporal information of the features, which are then input into a fully connected classification layer for classification and recognition. Summary of the Invention

[0005] To address the high computational complexity of existing multi-head attention mechanisms for processing one-dimensional time-series vibration signals, this application provides a gear fault identification method based on improved multi-head attention, which effectively reduces computational complexity.

[0006] One aspect of this application provides a gear fault identification method based on improved multi-head attention, comprising: acquiring vibration signals of gears under different fault states; preprocessing the vibration signals to obtain a standardized sample set; using a dual-path convolutional network to extract and fuse features from the samples in the standardized sample set to obtain a fused feature sequence; using an improved multi-head attention mechanism to enhance the features of the fused feature sequence to obtain an attention feature sequence; performing temporal feature encoding and global feature aggregation on the attention feature sequence to obtain global features; and identifying gear fault categories based on the global features.

[0007] Furthermore, the vibration signal is preprocessed, including: segmenting the vibration signal into multiple signal segments according to a preset length; normalizing each signal segment; assigning fault category labels to the normalized signal segments according to the fault type of the vibration signal; dividing the labeled signal segments into training and testing sets according to a preset ratio; and performing tensor transformation and dimension adjustment on the signal segments in the training and testing sets to obtain a standardized sample set.

[0008] Furthermore, each signal segment is normalized using the following formula: ;in, The original data matrix, and These represent the maximum and minimum values ​​for each feature dimension in the data matrix, respectively.

[0009] Furthermore, a dual-path convolutional network is used to extract and fuse features from samples in the standardized sample set, including: setting up a first convolutional path and a second convolutional path in parallel, wherein the kernel size of the first convolutional path is larger than the kernel size of the second convolutional path; extracting long-term features of the samples through the first convolutional path as the first feature; extracting short-term features of the samples through the second convolutional path as the second feature; and fusing the first feature and the second feature to obtain a fused feature sequence.

[0010] Furthermore, the kernel length of the first convolution path ranges from 10 to 20; the kernel length of the second convolution path ranges from 3 to 10.

[0011] Furthermore, an improved multi-head attention mechanism is used to enhance the features of the fused feature sequence, including: grouping the fused feature sequence into groups according to a preset number of attention heads; performing projection transformation on the fused feature sequence to generate query vectors, key vectors, and value vectors; wherein the key vectors and value vectors of multiple attention heads within the same group share the calculation; calculating attention weights based on the query vectors and key vectors; and weighting and aggregating the value vectors based on the attention weights to obtain the multi-head attention output.

[0012] Specifically, this application groups multiple attention heads, with each attention head within the same group sharing the projection calculations of key and value vectors, while maintaining its own independent query vector. This design breaks away from the traditional multi-head attention model where each head computes completely independently. While ensuring that different heads can focus on different feature patterns (through independent queries), it significantly reduces parameter redundancy by sharing key-value projections. Furthermore, generating a shared latent vector through a low-rank projection matrix serves as the basis for key-value calculations, further compressing the dimensionality of feature representations and reducing the computational load of subsequent projection transformations.

[0013] Furthermore, a projection transformation is performed on the fused feature sequence to generate a query vector, a key vector, and a value vector. This includes mapping the feature vector at each position in the fused feature sequence using a low-rank projection matrix to generate a shared latent vector. The shared latent vector is an intermediate representation vector obtained by dimensionality compression of the original feature vector using the low-rank projection matrix. This vector serves as the common basis for key-value calculations across multiple attention heads, reducing the feature dimensionality while preserving the main information of the original features, thus achieving parameter sharing among different attention heads.

[0014] Based on the position index of the feature vector in the fused feature sequence, calculate the corresponding position code; perform position fusion on the feature vector and shared latent vector of the fused feature sequence according to the position code, to obtain feature vectors with position information and shared latent vectors with position information respectively; set a query projection matrix for each attention head, and use the query projection matrix to perform a linear transformation on the feature vector with position information to generate a query vector; set a key projection matrix and a value projection matrix for each group, wherein multiple attention heads within the same group use the same key projection matrix and value projection matrix; use the key projection matrix to perform a linear transformation on the shared latent vector with position information to generate a key vector; use the value projection matrix to perform a linear transformation on the shared latent vector with position information to generate a value vector.

[0015] Specifically, this application introduces a low-rank projection matrix to map high-dimensional feature vectors into low-dimensional shared latent vectors, which serve as a unified input source for all key-value projections. Compared to traditional methods that directly perform multiple independent projections on the original high-dimensional features, this two-stage projection (first low-rank compression, then shared projection) significantly reduces the number of parameters.

[0016] Furthermore, although parameter sharing is employed, different attention heads can still focus on different feature patterns by maintaining the independence of query vectors and sharing key values ​​at the group level. Low-rank projection, by preserving the main feature components, maintains the ability to identify key features of gear faults while compressing parameters.

[0017] Furthermore, attention weights are calculated based on the query vector and key vector using the following formula:

[0018] ,in, Indicates the first t The query vector pairs the first s Attention weights for each key vector; Indicates the first i A position in the attention. t The corresponding query vector; Indicates the first i A position in the attention. s The corresponding key vector; Indicates the first i A position in the attention. j The corresponding key vector; Indicates the dimension of the key vector; m t, s, j = 1, 2, 3, …, m.

[0019] Furthermore, temporal feature encoding and global feature aggregation are performed on the attention feature sequence, including: constructing a bidirectional encoder containing a forward GRU network and a backward GRU network; inputting the attention feature sequence into the forward GRU network in chronological order, updating the hidden state through the gating mechanism of the forward GRU network to obtain a forward hidden state sequence; inputting the attention feature sequence into the backward GRU network in reverse chronological order, updating the hidden state through the gating mechanism of the backward GRU network to obtain a backward hidden state sequence; concatenating the forward and backward hidden states corresponding to each time step to obtain the bidirectional encoded features for the corresponding time step; combining the bidirectional encoded features of all time steps to obtain a bidirectional encoded sequence; and performing global average pooling on the bidirectional encoded sequence along the time dimension to obtain a global feature vector.

[0020] Another aspect of this application provides a gear fault identification system based on improved multi-head attention, comprising: a signal acquisition module for acquiring vibration signals of gears under different fault states; a preprocessing module for preprocessing the vibration signals to obtain a standardized sample set; a feature extraction module for using a dual-path convolutional network to extract and fuse features from samples in the standardized sample set to obtain a fused feature sequence; an attention enhancement module for using an improved multi-head attention mechanism to enhance the features of the fused feature sequence to obtain an attention feature sequence; a temporal coding module for performing temporal feature coding and global feature aggregation on the attention feature sequence to obtain global features; and a fault identification module for identifying gear fault categories based on the global features.

[0021] Compared to existing technologies, the advantages of this application are: (1) By introducing low-rank projection to generate shared latent vectors in the multi-head attention mechanism, and sharing the calculation of key vectors and value vectors among multiple attention heads in the same group, compared with the traditional multi-head attention mechanism that independently calculates all projection parameters for each attention head, this method reduces the number of parameters and the amount of computation to 1 / 3 to 1 / 2 of the original, and significantly reduces the computational resource consumption for model training and inference while maintaining the feature representation ability. (2) By using the parallel architecture of the dual-path convolutional network, long convolutional kernels (10-20) are used to capture the long-term periodic features of vibration signals, and short convolutional kernels (3-10) are used to capture short-term features such as transient impacts. After the two features are fused, the macroscopic evolution trend and microscopic abnormal patterns of gear faults can be preserved at the same time. Compared with the single-scale feature extraction method, it can more comprehensively characterize the features of different types of gear faults. Attached Figure Description Figure 1 This is a flowchart illustrating the gear fault identification method based on improved multi-head latent attention mentioned in the embodiments of this application; Figure 2 This is a schematic diagram of the gear fault classification results of a planetary gearbox according to an embodiment of this application; Figure 3 This is a radar chart comparing the performance of ablation experiments in this application with that of traditional multi-head attention systems. Detailed Implementation

[0022] The present application will now be described in detail with reference to the accompanying drawings and specific embodiments.

[0023] like Figure 1 As shown in the embodiments of this application, the gear fault identification method based on improved multi-head latent attention includes the following steps: S1. Acquire one-dimensional vibration and impact signals generated under various typical gear fault conditions and perform data preprocessing to provide data preparation for subsequent deep learning model training.

[0024] S101. Read all files storing vibration signals from the target directory. The file formats include common sensor signal storage formats (such as text, csv, mat, etc.), and each file corresponds to a gear fault state. Automatically identify the format based on the file extension and call the corresponding read function to complete the data loading.

[0025] S102. For each file's long sequence signal data, perform equal-length, non-overlapping segmentation according to a fixed length. Specifically, the j-th sample of the i-th type of fault is obtained using the following formula: The length n of each sample is set to 1024, meaning each sample contains 1024 data points to cover the complete impact cycle of gear failure and to meet the input size requirements of convolutional neural networks.

[0026] S103. Assign a corresponding fault category label to each sample Xi,j according to the file name of the data file. ,in , where i is the index of the fault category.

[0027] S104. Perform minimum-maximum normalization on the original vibration signal data to eliminate the influence of amplitude differences caused by different sensor gains and operating conditions. Linearly scale the data of each sample to the [0, 1] interval. The normalization formula is: ,in, The original data matrix, and These represent the maximum and minimum values ​​for each feature dimension in the data matrix, respectively.

[0028] S105. Divide the normalized samples into training and test sets in a 7:3 ratio. Use a stratified sampling strategy to ensure that the proportions of various fault samples in the training and test sets before and after the division are basically consistent with the original data.

[0029] S106. Convert the processed sample data (NumPy array format) into the tensor format required for deep learning. In this embodiment, the PyTorch framework is used. An additional channel dimension is added to accommodate the input requirements of the convolutional neural network. The specific steps are as follows: .

[0030] S2 expands the data dimension and extracts features from the one-dimensional vibration signal output by S1 based on dual-path convolution (DPC) to adapt to the subsequent multi-head attention head number division. Its core lies in extracting global and local features of the signal through two parallel convolutional paths, and achieving information complementarity and representation enhancement through feature fusion.

[0031] S201. Set up two parallel convolutional paths, each with a convolutional kernel of a specific size and a stride strategy, to extract macroscopic low-frequency features and microscopic high-frequency features of the signal respectively, and expand the signal dimension channels while extracting features.

[0032] S2011, Large Convolutional Kernel Path: In this embodiment, macroscopic feature extraction is achieved through a single convolution-pooling block. A single convolution-pooling block comprises two one-dimensional convolution operations and one pooling operation, specifically expressed as follows: The first convolutional layer uses a relatively large one-dimensional convolutional kernel with a length n1 of 18, a stride m1 of 2, and an output channel number O1 of 50. This aims to significantly expand the input features from one dimension to a high-dimensional space to capture diverse macroscopic features. The specific operation is as follows: .

[0033] The second convolutional layer uses a slightly smaller one-dimensional convolutional kernel with a length n² of 10, a stride of 2, and an output channel count O² of 30. It refines and moderately compresses the features from the previous layer. The specific steps are as follows: .

[0034] Pooling layer: One-dimensional max pooling is used, with both kernel length n3 and stride m3 set to 2 to compress the feature sequence length and enhance the translation invariance of the features. .

[0035] S2012, Small Convolutional Kernel Path: In this embodiment, macroscopic features are extracted through two convolution-pooling blocks. The length n4 of each small convolutional kernel is set to 6 to ensure consistency in local feature extraction; the stride m2 of the small convolutional kernel is set to 1 to enhance the ability to capture local features through denser sampling; the number of channels in the first layer is the same as O1, 50, to ensure feature alignment; the number of channels in the last layer is the same as O2, 30, to facilitate subsequent feature fusion. A gradually decreasing strategy is used for slow compression in the middle layers. Specifically: First convolutional layer: ; Second convolutional layer: ; First-level pooling: ; Third convolutional layer: ; Fourth convolutional layer: ; Second-level pooling: .

[0036] S202. After the convolution and pooling layers in each path, Adaptive Batch Normalization (AdaBN) is introduced to enhance the model's adaptability in different domains. Its calculation formula is as follows: Where x is the output feature of the convolutional layer. and ε represents the mean and variance estimates of the target domain data, respectively, and ε is a small constant used for numerical stability.

[0037] S203. After normalization, the adaptive activation function MetaAconC is used to dynamically enhance the nonlinear expressive power and generalization performance of the neural network. Its calculation formula is defined as: .

[0038] Where x is the input feature; p1 and p2 are learnable weight parameters dynamically generated through the network meta-learning mechanism; β is an adaptive scaling factor that controls the nonlinearity of the activation function; σ(·) is the Sigmoid function, and its formula is: Its domain is (-∞, +∞), and its range is (0, 1). When x = 0, the function value is 0.5. When x approaches -∞, the function value approaches 0. When x approaches +∞, the function value approaches 1.

[0039] S204. Fuse the feature P1 output by the macroscopic path with the feature P2 output by the microscopic path, and enhance the interaction between the features by element-wise multiplication. The specific operation is as follows: ,in This indicates element-wise multiplication. The feature dimensions remain unchanged before and after fusion, both being RB×C×L, where B, C, and L represent the batch size, number of channels, and feature sequence length, respectively.

[0040] S205. The fused 3D tensor features E are transposed and reshaped to convert them into the desired sequence input format X for subsequent multi-head attention. In subsequent attention calculations, to simplify the formula, the batch dimension B is omitted, so the dimension of X is... m=L is the sequence length, and d=C is the number of feature channels (i.e., model dimension).

[0041] S3. Based on the improved multi-head latent attention mechanism (IMLA), attention relationship modeling is performed on the feature sequence output in step S2. At the same time, the cache occupation in the inference stage is reduced by sharing parameters through low-rank projection, and sine and cosine coding is incorporated to enhance the model's perception of sequence position, solving the compatibility problem caused by the introduction of position information and the problem that the traditional multi-head latent attention mechanism cannot handle one-dimensional data.

[0042] S301. Grouping of Multi-Head Attention. The input sequence X is divided according to the number of attention heads h. The number of attention heads h is set to be equal to the model dimension d, and it is ensured that h is divisible by the number of groups g, i.e., h%g=0. In this embodiment, based on the compressed output dimension in S2, the number of attention heads h is set to 30, and the number of groups g is set to 5, i.e., each group contains 6 attention heads. Within each group, the query Q is calculated independently, while the key K and value V are shared.

[0043] S302. During the model training phase, perform a complete forward computation to learn all projection matrix parameters. The specific steps are as follows: S3021. Construct a shared latent vector. For each feature vector \(x_s\) corresponding to a position \(s\) in the input sequence \(X\) (\(s = 1, 2, \ldots, m\)), map it to a shared latent vector \(c_s\) through a low-rank projection matrix \(W_c\). The specific operation is as follows: . The specific operation formula is: ; where \(W_K\) and \(W_V\) are key-value projection matrices. Concatenate all and together and denote it as \(c_s\), and concatenate the corresponding projection matrices together and denote them as the low-rank projection matrix \(W_c\). The dimension of \(c_s\) is , \(d_k\) and \(d_v\) represent the query, key, and value dimensions of each head, and the calculation formula is: , then \(d_c < d\), so the change from \(x_s\) to \(c_s\) is a low-rank projection. In this embodiment, \(d_k = d_v = 30 / 30 = 1\), \(d_c = 5\) (1 + 1) = 10 < d = 30.

[0044] S3022. Inject position information.

[0045] The traditional multi-head latent attention mechanism (MLA) often uses Rotary Position Embedding (RoPE) to inject position information into the attention mechanism through a rotation matrix, and it often multiplies with the vector after projection. The original calculation methods of the query and key vectors containing RoPE are: ; where, is the rotation matrix corresponding to position \(s\).

[0046] However, when combined with low-rank projection, the traditional application method of RoPE will hinder matrix combination optimization in subsequent attention weight calculations, thus affecting the inference efficiency: ; Since the relative position rotation matrix is inserted between two projection matrices, and the value of the rotation matrix is related to the position, it is impossible to combine it with the projection matrix into a fixed matrix, resulting in the failure of low-rank projection sharing. To solve the above problem, the traditional MLA designs a hybrid-dimension RoPE design, which divides the query, key, and value of each attention head into a part without position encoding and a part with position encoding: , , . The dimension of the part with position encoding is \(d_r\), and its dimension is half of the original dimension: \(d_r = d_k / 2\). But if its dimension \(d_r\) is one-dimensional at this time, effective segmentation cannot be performed.

[0047] In contrast, the sine and cosine position coding PE(·) combines sine and cosine functions of different frequencies and adds them to the corresponding vectors without affecting the matrix merging optimization. Moreover, it can capture position information at different scales through the fundamental frequency f, making it more suitable for the multi-period and multi-frequency characteristics that may exist in fault signals.

[0048] Therefore, this embodiment uses sine and cosine positional encoding to address the incompatibility of rotational positional encoding with one-dimensional vectors in traditional multi-head latent attention mechanisms. The specific operation is as follows: The position coding vector is generated using the sine and cosine position coding functions PE(·). For position s, the sine and cosine position coding functions are defined as follows: Where f is a fixed fundamental frequency, which is 1024 in this embodiment; j is the dimension index. .

[0049] Location information is fused into the input feature xs and the shared latent vector cs to obtain the input feature with location information. and shared latent vectors .

[0050] S3023. Calculate the query, key, and value projection vectors for each attention head. For the i-th attention head (i=1, 2, ..., h), use an independent query projection vector. and the key-value projection matrix of the shared latent vector. , The calculation formula is as follows: ;in, , , Let represent the query, key, and value vector at the s-th position in the i-th header, respectively.

[0051] S3024. Calculate the multi-head attention output. First, calculate the attention weight matrix A(i) of the i-th attention head, where the attention weight of the element in the t-th row and s-th column is: Then the weighted sum of the outputs of the i-th attention head is: , where V(i) is a value matrix formed by stacking vs(i). Finally, the outputs of all attention heads are concatenated and fused using a linear transformation through the output projection matrix WO: .

[0052] S303. During the model inference phase, the merged projection matrix is ​​pre-computed, and the parameter cache from the training phase is utilized to avoid redundant calculations and reduce cache size. The specific operations are as follows: S3031. Initialization and Pre-computation of the Merge Matrix. During model loading, the query-key and query-value merged projection matrix for each attention head is pre-computed to eliminate redundant matrix multiplications during inference. The merged projection matrix As the projection matrix of Q, the query is now performed. AND key The dot product can be converted to: That is, replacing the original with CS. and This helps reduce cache usage.

[0053] S3032, Input sequence processing and positional encoding injection, are the same as the processing methods in S3021-S3022, to obtain input features xs' with positional information and shared latent vector cs'.

[0054] S3033. Calculate the transformed query vector based on the merge matrix. For the i-th attention head, the transformed query vector is calculated using the pre-computed merge matrix. Directly calculate the transformed query vector: S3034. Calculate attention weights. Calculate the attention weights using the shared latent vector with position from S3032 and the changed query vector from S3033. ,in The specific calculation is as follows: Among them, xt, cs, Since all of these operations have already been cached, the attention weight calculation during the inference phase can be performed quickly by caching only the position encoding (PE).

[0055] S3035, Calculate the attention head output and the final output. Use the pre-computed query-value merged projection matrix. Compute the equivalent value vector with the position-encoded shared latent vector cs' Then calculate the output of the i-th attention head: Finally, similar to the training phase, the outputs of all heads are concatenated and subjected to a linear transformation to obtain the final output.

[0056] S4. Based on the Bidirectional Gated Recurrent Unit (BiGRU), the attention optimization features output in step S3 are modeled temporally and sequentially, and compressed into a fixed-length global feature representation by global average pooling, providing a more discriminative feature vector for the final classification decision.

[0057] S401. Construct a bidirectional gated recurrent unit encoder. This encoder contains two parallel but oppositely directed GRU layers: a forward GRU for sequential processing and a backward GRU for reverse processing. Each GRU unit receives the current input xt and the hidden layer state from the previous time step. The specific update steps within the unit are as follows: S4011. At time step t, given the input vector And the hidden state of the previous step Calculate the reset gate: ,in, and These are the weight matrix and bias term of the reset gate, respectively. This indicates the concatenation of the current input with the previous hidden state.

[0058] S4012, Calculation Update Gate It controls the information of the previous hidden state. Pass to the hidden state at the current moment The proportion is calculated using the following formula: , where W and b are the weight matrix and bias term of the update gate, respectively.

[0059] S4013. Combine the previous hidden state (adjusted by the reset gate rt) with the current input xt to calculate the candidate hidden state: ,in, and are the weight matrix and bias term of the candidate hidden state, respectively, and tanh is the hyperbolic tangent activation function.

[0060] S4014, Finally, update the current hidden state ht through the update gate. right and The current hidden state is obtained by performing a linear combination: .

[0061] S402, Features of S3 Output The input is fed into a bidirectional GRU encoder for parallel processing. In this embodiment, the input feature dimension is set to 30 according to S3, and the hidden layer dimension is set to 64.

[0062] S4021, Forward GRU Encoding. Processes the sequence sequentially from front to back, capturing the dependencies between the current time step and past time steps. For the feature vector Ot at time step t, its forward hidden state... Based on the current input Ot and the previous hidden state The calculation shows that: .

[0063] S4022, Backward GRU encoding, processes the sequence backward from the beginning to capture the dependencies between the current and future time steps. Backward hidden state. The calculation formula is: ,in This is the hidden state for the next step.

[0064] S403. Concatenate the forward and backward hidden states of each time step along the feature dimension to form a bidirectional comprehensive code for that time step: Integrating these sequences yields the complete bidirectional encoded sequence: , where L is the sequence length, and its feature dimension becomes twice the original, i.e., 128.

[0065] S404. Perform global average pooling on the bidirectional encoded sequence H along the time dimension L, aggregating the feature information from all time steps into a fixed-length global feature vector: .

[0066] S5. Input the pooled feature vector into the fully connected classification layer, and its predicted output... The calculation formula is: Where Wfc and Bfc are the weight matrix and bias term of the classification layer, respectively. Finally, the probability distribution is converted into class labels: The argmax function returns the location of the maximum value in the vector, which is the final category index.

[0067] S6. Set up an early stopping mechanism to monitor the training process and prevent overfitting.

[0068] S601. At the start of training, initialize the early stop counter P=0 and set the maximum count value Pmax. In this embodiment, Pmax is set to 10.

[0069] S602. Train the model in batches according to steps S2-S5. After each training cycle, calculate the label smooth cross-entropy loss (LSR) of the validation set. This function can improve the generalization ability of the model. Its calculation formula is as follows: Where N is the number of label categories, and pi is the probability distribution predicted by the model. The i-th element, qi, is the smoothed true label, generated by the following formula: ; Where c represents the true category. The smoothing coefficient is set to 0.1 in this embodiment.

[0070] S7. Compare the current batch loss Linouk with the historical best loss Lbest.

[0071] S701. If Lnow < Lbest, update Lbest = Lnow, save the current model state, and reset the early stopping counter P to 0.

[0072] S702. If Lnow ≥ Lbest, then P + 1.

[0073] S703. When P ≥ Pmax, determine that the model performance no longer improves, terminate the training early, and load the model of the optimal batch saved previously.

[0074] S8. Save and export the optimal model after training as a.pt file, which contains the complete network structure definition and trained weight parameters. Finally, perform fault identification and output on the gear vibration fault signals in the target file.

[0075] This application also provides a gear fault identification system based on improved multi - head latent attention, including: Data pre - processing module: Data reading unit: Configured to automatically traverse and identify vibration signal data files with specific suffixes from a specified file directory, and call the corresponding library functions to load them.

[0076] Signal segmentation and annotation unit: Split each loaded long - sequence signal into non - overlapping segments of a fixed length to generate samples. At the same time, parse the fault type according to the file name or directory structure, and generate corresponding integer labels for each sample.

[0077] Normalization and partitioning unit: Apply the min - max normalization algorithm to the sample data matrix to linearly scale it. Subsequently, adopt a stratified sampling strategy to divide the entire data set into a training set and a test set according to a preset ratio.

[0078] Format conversion unit: Configured to convert the processed data into the tensor format of the deep learning framework, and appropriately increase the channel dimension to make its shape meet the input requirements of the convolutional network.

[0079] Dual - path convolutional feature extraction module: Large convolutional kernel path unit: Consists of a convolutional - pooling layer structure block, and after each convolutional and pooling layer, there is a batch normalization layer and a MetaAconC activation function.

[0080] Small convolutional kernel path unit: Consists of two convolutional - pooling layer structure blocks, and after each convolutional and pooling layer, there is a batch normalization layer and a MetaAconC activation function, and the number of input and output channels at the beginning and end is aligned with that of the large convolutional kernel path.

[0081] Feature fusion unit: Fuse the outputs of the large and small convolutional kernel path units by element - wise multiplication to generate fused features.

[0082] Improved multi-head potential attention module: Low-rank projection unit: projects each position vector of the input sequence into a low-dimensional shared latent vector, which is used to generate the keys and values ​​of all attention heads.

[0083] Position coding unit: It has a built-in sine and cosine position coding function, which generates a coding vector based on the sequence position and adds it to the input feature and the shared latent vector respectively.

[0084] Multi-head attention computation unit: Calculates the attention output of each head, and finally concatenates the outputs of all heads and performs a linear transformation through the output projection matrix to obtain the optimized feature sequence.

[0085] Bidirectional timing coding and compression module: Bidirectional gated loop unit: It uses a bidirectional channel to process time in forward and reverse order, and outputs an encoded sequence that integrates forward and backward information.

[0086] Global average pooling unit: Performs global average pooling on the sequence output by the bidirectional gated recurrent unit in the time dimension and compresses the sequence into a global feature vector.

[0087] Classification Decision Module: Classifier unit: Maps the global feature vector to the score of each category, uses the Softmax function to convert the score into a probability distribution, and obtains the final fault category prediction through the argmax operation.

[0088] Model training module: Loss calculation unit: The label smooth cross-entropy loss function is used to calculate the loss between the predicted value and the true label.

[0089] Optimization and Early Stop Control Unit: Monitors the validation set loss. If the loss does not decrease within a certain number of consecutive periods after reaching the maximum early stop value, early stop is triggered, and the model parameters with the lowest validation loss are saved as the final model.

[0090] This application also provides an electronic device, including one or more processors and a memory; One or more applications, wherein the one or more applications are stored in the memory and configured to be executed by the one or more processors, the one or more applications being configured to perform the gear fault identification method based on improved multi-head latent attention as described above.

[0091] This application also provides a non-transitory computer-readable storage medium having a computer program stored thereon, characterized in that, when the computer program is executed by a processor, it implements the gear fault identification method based on improved multi-head latent attention as described above.

[0092] Experimental verification: This experiment used gearbox data acquisition from Southeast University. The data files contained data on five different fault types—health, damage, broken teeth, tooth surface wear, and tooth root wear—under two operating conditions: a speed of 1200 rpm and a load of 0 N, and a speed of 1800 rpm and a load of 7.32 N. These data represented a total of 10 different fault categories. The data were processed according to the methods described in Examples S1-S8.

[0093] The confusion matrix of the experimental prediction results is shown in the figure below. Figure 2 As shown in the radar chart comparing the ablation test performance, see below. Figure 3 As shown in the figure (comparison between IMLA and the multi-head attention mechanism MHA), IMLA demonstrates significant improvements over MHA in all performance metrics. Specifically, IMLA reduces the number of training epochs from 52 to 40 (a reduction of 23.1%), training time from 113.3 seconds to 88.47 seconds (a reduction of 21.9%), and model accuracy from 92.67% to 95.67% (an increase of 3 percentage points). Regarding memory usage, IMLA reduces regular memory usage from 79.11 MB to 65.34 MB (a reduction of 17.4%), maximum memory usage from 1471.08 MB to 796.98 MB (a reduction of 45.8%), and cache usage from 1860 MB to 932 MB (a reduction of 49.9%). Experimental results demonstrate that IMLA significantly optimizes training efficiency and memory utilization while improving model accuracy, showcasing the practical significance of this application in gear fault diagnosis.

[0094] The foregoing illustrative description of the present application and its embodiments is not restrictive and can be implemented in other specific forms without departing from the spirit or essential characteristics of the present application. The accompanying drawings are only one embodiment of the present application, and the actual structure is not limited thereto. Therefore, if those skilled in the art are inspired by this description and design similar structures and embodiments without departing from the spirit of the present application, such designs should fall within the scope of protection of this application. Furthermore, the word "comprising" does not exclude other elements or steps, and the word "a" preceding an element does not exclude the inclusion of "a plurality" of that element. Terms such as "first," "second," etc., are used to indicate names and do not indicate any specific order.

Claims

1. A gear fault identification method based on improved multi-head attention, characterized in that, include: Obtain vibration signals of gears under different fault conditions; The vibration signals are preprocessed to obtain a standardized sample set; A dual-path convolutional network is used to extract and fuse features from samples in a standardized sample set to obtain a fused feature sequence. An improved multi-head attention mechanism is used to enhance the features of the fused feature sequence, resulting in an attention feature sequence. Temporal feature encoding and global feature aggregation are performed on the attention feature sequence to obtain global features; Gear fault categories are identified based on global features.

2. The gear fault identification method based on improved multi-head attention according to claim 1, characterized in that: Preprocessing of vibration signals includes: The vibration signal is divided into multiple signal segments according to a preset length; Normalize each signal segment; Based on the fault type of the vibration signal, a fault category label is set for the normalized signal segment; The labeled signal segments are divided into training and testing sets according to a preset ratio; Tensor transformation and dimension adjustment are performed on signal segments in the training and test sets to obtain a standardized sample set.

3. The gear fault identification method based on improved multi-head attention according to claim 2, characterized in that: Each signal segment is normalized using the following formula: Where Xorigin is the original data matrix, and Xmax and Xmin are the maximum and minimum values ​​of each feature dimension in the data matrix, respectively.

4. The gear fault identification method based on improved multi-head attention according to claim 2, characterized in that: A dual-path convolutional network is used to extract and fuse features from samples in a standardized sample set, including: Set up a first convolutional path and a second convolutional path in parallel, wherein the kernel size of the first convolutional path is larger than the kernel size of the second convolutional path; The long-term features of the samples are extracted through the first convolution path and used as the first feature; The short-timescale features of the samples are extracted through the second convolution path and used as the second feature; The first and second features are fused to obtain a fused feature sequence.

5. The gear fault identification method based on improved multi-head attention according to claim 4, characterized in that: The kernel length of the first convolution path ranges from 10 to 20; The kernel length of the second convolution path ranges from 3 to 10.

6. The gear fault identification method based on improved multi-head attention according to any one of claims 2 to 5, characterized in that: An improved multi-head attention mechanism is used to enhance the features of the fused feature sequence, including: The fused feature sequences are grouped according to a preset number of attention heads; The fused feature sequence is projected and transformed to generate a query vector, a key vector, and a value vector; the key vectors and value vectors of multiple attention heads within the same group are computed together. Calculate attention weights based on the query vector and key vector; The value vector is weighted and aggregated according to the attention weights to obtain the multi-head attention output.

7. The gear fault identification method based on improved multi-head attention according to claim 6, characterized in that: The fused feature sequence is projected and transformed to generate a query vector, a key vector, and a value vector, including: Based on the feature vector at each position in the fused feature sequence, a shared latent vector is generated by mapping through a low-rank projection matrix; Calculate the corresponding positional code based on the position index of the feature vector in the fused feature sequence; Based on the positional encoding, the feature vector and the shared latent vector of the fused feature sequence are fused to obtain the feature vector with positional information and the shared latent vector with positional information, respectively. Set the query projection matrix for each attention head, and use the query projection matrix to perform a linear transformation on the feature vector with positional information to generate the query vector; Set the key projection matrix and value projection matrix for each group, where multiple attention heads within the same group use the same key projection matrix and value projection matrix; A linear transformation is performed on the shared latent vector with location information using the key projection matrix to generate a key vector. A value vector is generated by linearly transforming a shared latent vector with location information using a value projection matrix.

8. The gear fault identification method based on improved multi-head attention according to claim 6, characterized in that: The attention weights are calculated based on the query vector and key vector using the following formula: ,in, Indicates the first t The query vector pairs the first s Attention weights for each key vector; Indicates the first i A position in the attention. t The corresponding query vector; Indicates the first i A position in the attention. s The corresponding key vector; Indicates the first i A position in the attention. j The corresponding key vector; Indicates the dimension of the key vector; m The length of the key-value sequence.

9. The gear fault identification method based on improved multi-head attention according to claim 6, characterized in that: Temporal feature encoding and global feature aggregation are performed on the attention feature sequence, including: Construct a bidirectional encoder that includes a forward GRU network and a backward GRU network; The attention feature sequence is input into the feedforward GRU network in chronological order, and the hidden state is updated through the gating mechanism of the feedforward GRU network to obtain the feedforward hidden state sequence. The attention feature sequence is input into the backward GRU network in reverse chronological order, and the hidden state is updated through the gating mechanism of the backward GRU network to obtain the backward hidden state sequence. The forward hidden state and backward hidden state corresponding to each time step are concatenated to obtain the bidirectional encoded features of the corresponding time step. By combining the bidirectional encoded features of all time steps, a bidirectional encoded sequence is obtained; The bidirectional encoded sequence is globally averaged and pooled along the time dimension to obtain the global feature vector.

10. A gear fault identification system based on improved multi-head attention, used to implement the method according to any one of claims 1 to 9, characterized in that, include: The signal acquisition module is used to acquire vibration signals of the gear under different fault conditions; The preprocessing module is used to preprocess the vibration signal to obtain a standardized sample set; The feature extraction module is used to extract and fuse features from samples in the standardized sample set using a dual-path convolutional network to obtain a fused feature sequence. The attention enhancement module is used to enhance the fused feature sequence using an improved multi-head attention mechanism to obtain an attention feature sequence. The temporal coding module is used to perform temporal feature coding and global feature aggregation on the attention feature sequence to obtain global features; The fault identification module is used to identify gear fault categories based on the global features.

Images