Fault diagnosis method and system for improving vit architecture based on hierarchical attention mechanism

Abstract

The invention relates to the technical field of fault diagnosis, in particular to a fault diagnosis method and system for improving vit architecture based on a hierarchical attention mechanism. The fault diagnosis method comprises the steps of extracting a sample signal for preprocessing to obtain a spectrogram divided into a training set and a test set, constructing a feature enhancement module by combining two convolution layers, a pooling layer and a channel attention module, adding a patch partition layer, and applying a multi-stage hierarchical multi-head attention mechanism to obtain a fault diagnosis result. Hierarchical encoders with different layer numbers, different head numbers and different embedding dimensions are configured among the stages, an improved vit architecture model is built, a bearing fault diagnosis model is obtained through training of a spectrogram in a training set, the spectrogram in a test set is input into the trained bearing fault diagnosis model, and a fault diagnosis result is obtained. And outputting to obtain a corresponding fault diagnosis result. According to the method, multi-dimensional features can be efficiently extracted, the diagnosis effect is enhanced, and the defects of traditional fault diagnosis are effectively overcome.

Description

Technical Field

[0001] The present invention relates to the technical field of fault diagnosis, and particularly to a fault diagnosis method and system based on an improved vit architecture with a hierarchical attention mechanism. Background Art

[0002] With the continuous growth of the operating efficiency and automation requirements of industrial equipment, the reliability and fault diagnosis of equipment bearings have become key research fields. Traditional bearing fault diagnosis methods usually rely on manual experience and rule-based diagnosis means, such as spectrum analysis and signal filtering. Although these methods perform well in simple bearing fault modes, they are difficult to cope with the variability and complexity of different signal characteristics when facing diverse and complex bearing fault types. For example, traditional methods handle non-stationary signals, signals with strong noise interference, and situations where multiple fault characteristics coexist poorly, resulting in insufficient accuracy and robustness in bearing fault detection. In addition, traditional signal processing methods rely on specific professional domain knowledge and are less adaptable when dealing with large-scale data and signals of new equipment.

[0003] Currently, with the rapid development of deep learning technology, new ideas have been provided for the feature extraction of complex signals and fault mode recognition. Especially the introduction of the ViT model (Vision Transformer) and the self-attention mechanism has enabled models based on image and sequence data to exhibit excellent performance in industrial fault diagnosis. The ViT model is a computer vision model based on the Transformer architecture, which can efficiently capture the long-range dependence relationships and global features of signals through the multi-head self-attention mechanism, adapt to different levels of feature information fusion, and thus has strong feature extraction capabilities and adaptability. However, existing bearing fault diagnosis models are mostly designed for specific signal types or fault modes, often lacking the ability to adapt to complex signal characteristics and being difficult to achieve excellent results in multi-level information fusion. Summary of the Invention

[0004] The technical problem to be solved in the embodiments of the present invention is to provide a fault diagnosis method and system based on an improved vit architecture with a hierarchical attention mechanism to solve the above-mentioned deficiencies in bearing fault diagnosis in the prior art.

[0005] The present invention discloses a fault diagnosis method based on an improved vit architecture with a hierarchical attention mechanism, including: Obtain the vibration data of the bearings of the target device, extract sample signals from the vibration data and divide them into a training set and a test set, and convert the sample signals of the training set and the test set into spectrograms respectively; Based on two layers of lightweight convolutional layers and a pooling layer as the network foundation, combine a channel attention module to construct a feature enhancement module; Add a patch partition layer after the feature enhancement module, and apply a multi-stage hierarchical multi-head attention mechanism after the patch partition layer. Hierarchical encoders with different numbers of layers, different numbers of heads, and different embedding dimensions are configured between each stage; Build an improved vit architecture model with the feature enhancement module, patch partition layer, and hierarchical multi-head attention mechanism. Input the spectrograms in the training set into the vit architecture model for model training to obtain a bearing fault diagnosis model; Input the spectrograms in the test set into the trained bearing fault diagnosis model, and output the corresponding fault diagnosis results.

[0006] Optionally, the conversion of the sample signals in the training set and test set into spectrograms includes: Screen out the horizontal vibration signals from the vibration data, extract sample signals from the horizontal vibration signals with a preset step size and overlap rate, and randomly divide the extracted sample signals into a training set and a test set; Extract signal features from the divided sample signals at different frequencies and time scales through continuous wavelet transform, and convert them into spectrograms with preset pixel values. The functional expression of the continuous wavelet transform is:

[0007] In the formula, represents the mother wavelet function, s represents the scale factor, τ represents the time translation factor, represents the time variable, represents the mother wavelet function after stretching and translation.

[0008] Optionally, the fault diagnosis method further includes a method of processing the input spectrogram using two layers of lightweight convolutional layers, including: Use the first convolutional layer to perform a convolution operation on the input spectrogram with a preset convolutional kernel and step size. While keeping the size of the output feature map consistent with the input spectrogram, expand the output channels to the number of the first channels; Add a batch normalization layer after the first convolutional layer to normalize the eigenvalues of each output channel, and introduce a Relu activation function after the batch normalization layer to suppress negative activation of each normalized eigenvalue to obtain the first type of feature map after local feature extraction; Use the second convolutional layer to perform a convolution operation on the input first type of feature map with the same convolutional kernel and step size. While keeping the size of the output feature map consistent with the input feature map, expand the output channels to the number of the second channels, and output the second type of feature map after local feature extraction. The number of the second channels is greater than the number of the first channels; A batch normalization layer is added after the second convolutional layer to normalize the eigenvalues of each output channel, and a Relu activation function is introduced after the batch normalization layer to suppress negative activation of each normalized eigenvalue, obtaining the second type of feature map after local feature extraction.

[0009] Optionally, the fault diagnosis method further includes a method for performing convolution on a given position on the input spectrogram or feature map, including: According to the convolutional layer used, traverse each position of the convolutional kernel, and use padding to add zero padding to the edges of the input spectrogram or feature map; Slide the convolutional kernel on the input spectrogram or feature map to the given position, and calculate the output feature map of the given position. The functional expression for the output of the feature map at the given position is:

[0010] In the formula, k represents the output channel index, i represents the index of the corresponding output feature map in the height direction, j represents the index of the corresponding output feature map in the width direction, K represents the convolutional kernel size, u represents the position index of the convolutional kernel in the height direction, v represents the position index of the convolutional kernel in the width direction, represents the input spectrogram or feature map, c represents the input channel number, and p represents the offset generated by using padding, represents the given position on the input spectrogram or feature map The feature map output after convolution, represents the convolutional kernel weight, represents the bias, represents the number of channels of the input spectrogram or feature map.

[0011] Optionally, the fault diagnosis method further includes a method for processing the input second type of feature map using a pooling layer, including: Use the pooling layer to perform pixel filling on the input second type of feature map with a pooling kernel of a preset size; Slide the pooling window of the preset size on the filled second type of feature map according to the preset stride, and extract the maximum eigenvalue within the pooling window after each slide, and output the third type of feature map after pooling. The functional expression for outputting the third type of feature map after pooling is:

[0012] In the formula, represents the third type of feature map output after the pooling operation at the given position on the second type of feature map, represents the filled second type of feature map, and m and n represent the relative offsets within the pooling window.

[0013] Optionally, the fault diagnosis method further includes a method of processing the input third type of feature map using a channel attention module, including: Extract the features of each channel of the third type of feature map, introduce an adaptive weight assignment layer on each channel, and the adaptive weight assignment layer dynamically adjusts the weights of each channel according to the feature responses of the channels; At the same time, use global average pooling and global max pooling to capture the statistical information of each channel feature, send the obtained statistical information to a shared multi-layer perceptron, introduce a first fully connected layer to map the statistical information of global average pooling to the weight value of each channel, and introduce a second fully connected layer to map the statistical information of global max pooling to the weight value of each channel; Dynamically adjust the weights of each channel feature response according to the mapped weight values, and perform weighted combination on the features of each channel through the mapped weight values to output a channel attention map. The function expression for outputting the channel attention map is:

[0014] In the formula, Channel attention module function, Represents the input data, Represents the first fully connected layer, Represents the second fully connected layer, Is Function processing, Represents the weight of the first fully connected layer, Represents the result of global average pooling, Represents the weight of the second fully connected layer, Represents the result of global max pooling.

[0015] Optionally, the fault diagnosis method further includes a method of processing the input channel attention map using a patch partition layer, including: Set the partition window size of the patch partition layer, and by sliding the partition window size on the channel attention map, cut the channel attention map into multiple non-overlapping image blocks along the spatial dimension, and output the partitioned fourth type of feature map.

[0016] Optionally, the fault diagnosis method further includes a method of processing the input fourth type of feature map using a hierarchical multi-head attention mechanism, including: Add a linear embedding layer to perform linear projection on the input fourth type of feature map, map it to an embedding vector of a fixed dimension, then perform spatial compression and expand the spatial dimension, adjust the dimension order to obtain an embedding sequence, use a hierarchical encoder with a preset number of layers, number of heads, and embedding dimension to preliminarily extract features from the embedding sequence, and apply downsampling to increase the feature dimension, and then output a first-stage feature map; Add a linear embedding layer to perform spatial compression on the input first-stage feature map and flatten it into a sequence. Use a hierarchical encoder with a preset number of layers, number of heads, and embedding dimension to further extract features from the flattened sequence. The number of heads and embedding dimension of the second-stage hierarchical encoder are greater than those of the first-stage hierarchical encoder. After applying downsampling to increase the feature dimension, output the second-stage feature map; Add a linear embedding layer to perform spatial compression on the input second-stage feature map and flatten it into a sequence. Use a hierarchical encoder with a preset number of layers, number of heads, and embedding dimension to deeply extract features from the flattened sequence. The number of layers, number of heads, and embedding dimension of the third-stage hierarchical encoder are greater than those of the second-stage hierarchical encoder. After applying downsampling to increase the feature dimension, output the third-stage feature map; Add a linear embedding layer to perform spatial compression on the input third-stage feature map and flatten it into a sequence. Use a hierarchical encoder with a preset number of layers, number of heads, and embedding dimension to establish global relationships for the input sequence and output a high-level feature sequence after fusing the first three stages. The number of heads and embedding dimension of the fourth-stage hierarchical encoder are greater than those of the third-stage hierarchical encoder; Normalize all the output high-level feature sequences, and perform global average pooling on the normalized high-level feature sequences along the sequence length dimension. The functional expression of the global average pooling is:

[0017] In the formula, represents the result after global pooling, represents the length after the feature map is flattened, represents the i-th feature value of the input; Use a fully connected layer to map the high-level feature sequence after global average pooling to a preset output dimension, and output the corresponding fault diagnosis result. The functional expression of the mapping of the high-level feature sequence is:

[0018] In the formula, the result of the mapping output after global average pooling, the weight matrix in the learning parameters, the bias in the learning parameters.

[0019] Optionally, the fault diagnosis method further includes a method for extracting features from the sequence using a hierarchical encoder, including: Add layer normalization to normalize the input sequence, calculate self-attention for the normalized sequence through the multi-head attention mechanism, add the output of the self-attention to the original input sequence through residual connection, and output the first feature with the same dimension as the input sequence. The functional expression of the output of the first feature is:

[0020]

[0021] In the formula, represents the first feature of the output, X represents the input sequence, represents the normalized sequence, represents the output vector calculated by the attention function for the normalized sequence, represents the output vector calculated by the attention function, Q represents the query vector, K represents the key vector, and V represents the value vector, represents a sparse vector of the same size as Q, D represents the key vector dimension, and softmax represents the classification function; Add layer normalization to normalize the first feature of the input, apply a feed-forward layer to perform a non-linear transformation on the normalized first feature through a multi-layer perceptron, and add the output of the feed-forward layer to the original input first feature through a residual connection to output a second feature with the same dimension as the input first feature. The functional expression of the output of the second feature is:

[0022] In the formula, represents the second feature of the output, represents the normalized first feature, represents the output vector of the non-linear transformation of the normalized first feature through the multi-layer perceptron; Reference the patch merging layer to splice adjacent output vectors in the local area to obtain a composite feature vector, normalize the composite feature vector along the channel dimension, and use a fully connected layer to compress the normalized composite feature vector to the target number of channels to output the final feature map. The functional expression of the output of the final feature map is:

[0023] In the formula, represents the output final feature map, represents the output vector, represents the normalized composite feature vector, represents the compression of the composite feature vector to the target number of channels.

[0024] The present invention also discloses a fault diagnosis system, including: A data preprocessing module for acquiring vibration data of the bearing of the target device, extracting sample signals from the vibration data and dividing them into a training set and a test set, and converting the sample signals of the training set and the test set into spectrograms respectively; The basic network construction module is used to build a feature enhancement module by combining a channel attention module with a network foundation consisting of two lightweight convolutional layers and a pooling layer; The self-attention mechanism construction module is used to add a patch partition layer after the feature enhancement module and apply a multi-stage hierarchical multi-head attention mechanism after the patch partition layer. Hierarchical encoders with different numbers of layers, different numbers of heads, and different embedding dimensions are configured between each stage; The model establishment and training module is used to build an improved vit architecture model with the feature enhancement module, patch partition layer, and hierarchical multi-head attention mechanism, input the spectrogram in the training set into the vit architecture model for model training, and obtain a bearing fault diagnosis model; The fault diagnosis module is used to input the spectrogram in the test set into the trained bearing fault diagnosis model and output the corresponding fault diagnosis result.

[0025] Compared with the prior art, the beneficial effects of the fault diagnosis method and system based on the improved vit architecture with hierarchical attention mechanism provided by the embodiments of the present invention are as follows: By integrating convolution pooling, channel attention mechanism, and hierarchical multi-head attention mechanism, it can efficiently extract multi-dimensional features from the vibration data of the bearings of the target device, enhance the diagnostic effect using the multi-dimensional features of the vibration data, effectively overcome the deficiencies of traditional fault diagnosis, and can achieve high-precision and strong-robustness fault diagnosis for different devices and different fault modes. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] The technical solutions of the present invention will be further described in detail below in conjunction with the drawings and embodiments. In the drawings: Figure 1 It is a schematic block diagram of the steps of the fault diagnosis method provided by the embodiment of the present invention; Figure 2 It is a schematic flow chart of the fault diagnosis method provided by the embodiment of the present invention; Figure 3 It is a schematic flow chart of the process of the bearing fault diagnosis model processing the input spectrogram provided by the embodiment of the present invention; Figure 4 It is a schematic flow chart of the process of the feature enhancement module processing the input spectrogram provided by the embodiment of the present invention; Figure 5 It is a schematic structural diagram of the channel attention mechanism layer provided by the embodiment of the present invention; Figure 6 It is a schematic structural diagram of the hierarchical multi-head attention mechanism provided by the embodiment of the present invention. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0027] It should be noted that, without conflict, the embodiments in the present application and the features in the embodiments can be combined with each other. Now, in conjunction with the accompanying drawings, the preferred embodiments of the present invention will be described in detail.

[0028] The present invention discloses a fault diagnosis method for improving the vit architecture based on a hierarchical attention mechanism, as Figure 1 and Figure 2 shown, including: Obtain the vibration data of the bearings of the target device, extract sample signals from the vibration data and divide them into a training set and a test set, and convert the sample signals of the training set and the test set into spectrograms respectively; Take two layers of lightweight convolutional layers and a pooling layer as the network basis, and combine a channel attention module to construct a feature enhancement module; Add a patch partition layer after the feature enhancement module, and apply a multi-stage hierarchical multi-head attention mechanism after the patch partition layer. Hierarchical encoders with different numbers of layers, different numbers of heads, and different embedding dimensions are configured between each stage; Build an improved vit architecture model with a feature enhancement module, a patch partition layer, and a hierarchical multi-head attention mechanism. Input the spectrograms in the training set into the vit architecture model for model training to obtain a bearing fault diagnosis model; Input the spectrograms in the test set into the trained bearing fault diagnosis model, and output the corresponding fault diagnosis results.

[0029] Through the implementation of the above embodiments of the fault diagnosis method, combined with convolutional pooling, channel attention mechanism, and hierarchical multi-head attention mechanism, multi-dimensional feature information in the signal can be comprehensively captured. The convolutional layer is responsible for capturing and expanding the local features of the image, while the pooling layer ensures the effective operation of subsequent modules by screening and enhancing significant features while maintaining the spatial size unchanged. The channel attention mechanism highlights important spectral information by adaptively weighting the features of each channel. The hierarchical multi-head attention mechanism can analyze the signal from different perspectives, perform global feature modeling, and perfectly extract the features of the fault signal. By using a multi-stage hierarchical encoder configured with different numbers of layers, different numbers of heads, and different embedding dimensions, multi-dimensional information of the signal can be processed in parallel, enabling the bearing fault diagnosis model to maintain high robustness and adaptability in the face of complex signal interference, noise, and multiple fault modes, and helping the bearing fault diagnosis model to gradually learn and strengthen the understanding of the deep features in the bearing signal, ensuring stable output of accurate fault diagnosis results under various working conditions. Therefore, by combining convolutional pooling, channel attention, and multi-head self-attention mechanisms, the bearing fault diagnosis model can effectively fuse different features, automatically assign appropriate weights to each part of different signals, and this ability of multi-level feature fusion helps to extract the most discriminative features and improve the accuracy of bearing fault diagnosis. Further, the sample signals of the training set and the test set are respectively converted into spectrograms, including: Screen out the horizontal vibration signals from the vibration data, extract sample signals from the horizontal vibration signals at a preset step size and overlap rate, and randomly divide the extracted sample signals into a training set and a test set; Extract signal features from different frequencies and time scales for the divided sample signals through continuous wavelet transform, and convert them into spectrograms with preset pixel values. The functional expression of the continuous wavelet transform is:

[0030] In the formula, represents the mother wavelet function, s represents the scale factor, τ represents the time translation factor, represents the time variable, represents the mother wavelet function after scaling and translation.

[0031] By implementing the embodiments of the above fault diagnosis method, it is preferable to collect vibration data through sensors, extract the horizontal time-domain vibration signal with a step size of 1024 and an overlap rate of 0.5, and then randomly divide it into a training set, a test set, and a validation set, which helps to evaluate the generalization ability of the model and avoid overfitting. Among them, the training set is used to train the model to learn the characteristics of the data; the validation set is used to adjust the model parameters such as the learning rate and regularization parameters during the training process; the test set is used to input into the trained bearing fault diagnosis model for fault diagnosis output. By performing continuous wavelet transform on the horizontal time-domain vibration signal, the signal can be analyzed at different time scales to provide information about the variation of different frequency components in the signal over time, so that the time-frequency characteristics of the signal can be extracted. Then, the transformed eigenvalue is normalized to a preset range value, which helps to convert eigenvalues of different magnitudes to the same scale, making them comparable in subsequent processing. Preferably, the sample signal can be converted into a spectrogram of 224×224 pixels to optimize the visual effect and make it more suitable for subsequent fault analysis.

[0032] Furthermore, in combination with Figure 3 and Figure 4 as shown, the fault diagnosis method further includes a method of processing the input spectrogram using two layers of lightweight convolutional layers, including: Using the first convolutional layer to perform a convolution operation on the input spectrogram with a preset convolution kernel and step size, while keeping the size of the output feature map the same as that of the input spectrogram, expanding the output channels to the number of the first channels; Adding a batch normalization layer after the first convolutional layer to normalize the eigenvalues of each output channel, and introducing a Relu activation function after the batch normalization layer to suppress negative activation of each normalized eigenvalue, obtaining the first type of feature map after local feature extraction; Using the second convolutional layer to perform a convolution operation on the input first type of feature map with the same convolution kernel and step size, while keeping the size of the output feature map the same as that of the input feature map, expanding the output channels to the number of the second channels, and outputting the second type of feature map after local feature extraction, where the number of the second channels is greater than the number of the first channels; Adding a batch normalization layer after the second convolutional layer to normalize the eigenvalues of each output channel, and introducing a Relu activation function after the batch normalization layer to suppress negative activation of each normalized eigenvalue, obtaining the second type of feature map after local feature extraction.

[0033] Furthermore, the fault diagnosis method further includes a method of performing convolution on a given position on the input spectrogram or feature map, including: According to the convolutional layer used, traverse each position of the convolution kernel, and add zero padding to the edge of the input spectrogram or feature map using padding; Slide the convolution kernel over the input spectrogram or feature map to a given position, and calculate the feature map at the given position. The functional expression for the output of the feature map at the given position is:

[0034] In the formula, k represents the output channel index, i represents the index of the corresponding output feature map in the height direction, j represents the index of the corresponding output feature map in the width direction, K represents the convolution kernel size, u represents the position index of the convolution kernel in the height direction, v represents the position index of the convolution kernel in the width direction, represents the input spectrogram or feature map, c represents the input channel number, and p represents the offset generated by using padding. represents a given position on the input spectrogram or feature map The feature map output after convolution, represents the convolution kernel weight, represents the bias, represents the number of channels of the input spectrogram or feature map.

[0035] Through the implementation of the above embodiments of the fault diagnosis method, the convolutional layer is mainly responsible for capturing and expanding the local features of the image. Both convolutional layers perform convolution with a convolution kernel of 3 and a stride of 1. While extracting local information, it ensures that the size of the feature map after convolution does not change. For example: the first convolutional layer can expand the channels of the feature map to 64, and the second convolutional layer can expand the channels of the feature map to 128. Using two lightweight convolutional layers can not only extract sufficient local fault features, but also provide more channel dimensions for the subsequent channel attention module, and finally form a feature map after local feature extraction. Among them, padding (zero padding) is a commonly used technique in computer science and data processing. By supplementing additional values (usually zeros) at the edges of the data, it adjusts the size or shape of the data, and is widely used in fields such as deep learning, image processing, and signal processing.

[0036] For the two-dimensional convolutional layer in the embodiments of the present invention, its specific structure refers to Figure 3 Assume that the size of the input spectrogram or feature map is and the convolution kernel weight is Here The bias is (the bias is 0 in the bearing fault diagnosis model of the present invention). The feature map at the given position can be calculated through the above functional expression for the output of the feature map at the given position. Since the stride = 7 and the extended p = 0, each output element corresponds to a non-overlapping region in the input of Finally, the height and width of the output feature map are both 32 (224 / 7).

[0037] Furthermore, the fault diagnosis method further includes a method of processing the input second type of feature map using a pooling layer, including: Using a pooling layer to perform pixel padding on the input second type of feature map with a pooling kernel of a preset size; Sliding a pooling window of a preset size over the padded second type of feature map at a preset stride, and extracting the maximum feature value within the pooling window after each slide, and outputting a third type of feature map after pooling. The functional expression for outputting the third type of feature map after pooling is:

[0038] In the formula, represents the third type of feature map output after pooling operation at a given position on the second type of feature map, represents the padded second type of feature map, and m and n represent the relative offsets within the pooling window.

[0039] Through the implementation of the above embodiments of the fault diagnosis method, the pooling layer is used to ensure the effective operation of subsequent modules by screening and enhancing significant features while keeping the spatial dimensions unchanged. The pooling layer is preferably a pooling window with a size of 2 and a stride of 1, and after pooling, the pixels of the feature map will be reduced by one column and one row, and then the feature map is restored by filling with blank pixels to make the size of the feature map input into the pooling layer the same as the output size. In the embodiments of the present invention, a custom pooling layer is used, and the key lies in performing additional padding on the right side and the lower side of the input tensor, so that the output size after max pooling is the same as the input. For example: first, pad the input second type of feature map, that is, fill 1 pixel on the right side and the lower side of the second type of feature map, and then slide a window with a stride of 1, take the maximum value of the values within each window, and output the third type of feature map after pooling. This design ensures that the output size is the same as the original input, and at the same time can select the most significant activation value in the local area.

[0040] The design of using two layers of lightweight convolutional layers and one layer of pooling layer enables the bearing fault diagnosis model to establish a rich and robust feature representation at an early stage, providing a solid foundation for the overall task.

[0041] Furthermore, as shown in combination with Figure 3 and Figure 5 , the fault diagnosis method further includes a method of processing the input third type of feature map using a channel attention module, including: Extracting the features of each channel of the third type of feature map, introducing an adaptive weight assignment layer on each channel, and the adaptive weight assignment layer dynamically adjusts the weights of each channel according to the feature responses of the channels; Simultaneously use global average pooling and global max pooling to capture the statistical information of the features of each channel, send the obtained statistical information to a shared multi-layer perceptron, introduce a first fully connected layer to map the statistical information of global average pooling to the weight value of each channel, and introduce a second fully connected layer to map the statistical information of global max pooling to the weight value of each channel; Dynamically adjust the weight of the feature response of each channel according to the mapped weight value, and perform weighted combination on the features of each channel through the mapped weight value to output a channel attention map. The function expression for outputting the channel attention map is:

[0042] In the formula, Channel attention module function, represents the input data, represents the first fully connected layer, represents the second fully connected layer, is function processing, represents the weight of the first fully connected layer, represents the result of global average pooling, represents the weight of the second fully connected layer, represents the result of global max pooling.

[0043] Through the implementation of the above embodiments of the fault diagnosis method, use the channel attention module to process the third type of feature map, use the channel attention module to extract the features of each channel on the third type of feature map, and ensure that the signal patterns of each channel can be independently analyzed. Dynamically adjust the weight according to the feature response of different channels to highlight the key spectrum information. Introduce an adaptive weight assignment layer on each channel, use global average pooling or global max pooling to obtain channel statistical information, and send it to a shared multi-layer perceptron (MLP). The high-dimensional statistical information can be converted into weight values that can be used for weighted combination through a fully connected layer. So that the bearing fault diagnosis model can understand the overall trend and importance of the features of each channel, and can dynamically adjust the weight of the feature response of each channel according to the mapped weight value, so that the features of some channels are given higher weights, and thus show the importance of the features of each channel in the final output through the channel attention map.

[0044] Furthermore, the fault diagnosis method also includes a method of using a patch partition layer to process the input channel attention map, including: Set the partition window size of the patch partition layer, and by sliding the partition window size on the channel attention map, cut the channel attention map into multiple non-overlapping image blocks along the spatial dimension, and output the partitioned fourth type of feature map.

[0045] By implementing the embodiments of the above-mentioned fault diagnosis method, the patch partition layer is used to convert the channel attention map into a series of non-overlapping image patches, and the feature map is divided into grid-like patches in the spatial dimension, and each patch is processed independently. The divided feature map is used as the input for the subsequent stage and may be used for feature extraction at different levels. Thus, partitioning is performed on the weighted channel attention map to strengthen the extraction of local information of important channels.

[0046] Furthermore, as shown in Figure 3 and Figure 6 , the fault diagnosis method further includes a method of processing the input fourth type of feature map using a hierarchical multi-head attention mechanism, including: Adding a linear embedding layer to perform linear projection on the input fourth type of feature map, mapping it to an embedding vector of a fixed dimension, then performing spatial compression and expanding the spatial dimension, adjusting the dimension order to obtain an embedding sequence, and using a hierarchical encoder with a preset number of layers, number of heads, and embedding dimension to initially extract features from the embedding sequence, and applying downsampling to increase the feature dimension, and then outputting a first-stage feature map; Adding a linear embedding layer to perform spatial compression on the input first-stage feature map and flatten it into a sequence, using a hierarchical encoder with a preset number of layers, number of heads, and embedding dimension to further extract features from the flattened sequence. The number of heads and embedding dimension of the second-stage hierarchical encoder are greater than those of the first-stage hierarchical encoder, and applying downsampling to increase the feature dimension, and then outputting a second-stage feature map; Adding a linear embedding layer to perform spatial compression on the input second-stage feature map and flatten it into a sequence, using a hierarchical encoder with a preset number of layers, number of heads, and embedding dimension to deeply extract features from the flattened sequence. The number of layers, number of heads, and embedding dimension of the third-stage hierarchical encoder are greater than those of the second-stage hierarchical encoder, and applying downsampling to increase the feature dimension, and then outputting a third-stage feature map; Adding a linear embedding layer to perform spatial compression on the input third-stage feature map and flatten it into a sequence, using a hierarchical encoder with a preset number of layers, number of heads, and embedding dimension to establish a global relationship for the input sequence, and outputting a high-level feature sequence after fusing the first three stages. The number of heads and embedding dimension of the fourth-stage hierarchical encoder are greater than those of the third-stage hierarchical encoder; Performing normalization processing on all the output high-level feature sequences, and performing global average pooling on the normalized high-level feature sequences along the sequence length dimension. The function expression of the global average pooling is:

[0047] In the formula, represents the result after global pooling, represents the length after the feature map is flattened, represents the i-th eigenvalue of the input; Use a fully connected layer to map the high-level feature sequence after global average pooling to a preset output dimension, and output the corresponding fault diagnosis result. The functional expression for the mapping of the high-level feature sequence is:

[0048] In the formula, the result of mapping and output after global average pooling, the weight matrix in the learning parameters, the bias in the learning parameters.

[0049] Through the implementation of the above embodiments of the fault diagnosis method, the linear embedding layer is used to multiply the height and width of the feature map to form sequence data that can be processed by the ViT encoder. In the first stage, the height and width of the feature map are divided by 2 again, that is, 16x16, and then passed into the linear embedding layer to transform the feature map into a series of partitioned embedding representations, and the output spatial dimensions (height and width) are unfolded and converted into a sequence form, and the dimension order of the data is adjusted to the format of batch size (BatchSize, B), sequence length (Number of Elements, N), and feature dimension (Channels, C). Preferably, B = 32, N = 32×32 = 1024, and C = 192.

[0050] The number of heads of the first-stage hierarchical encoder is 3, the number of layers is 2, and the embedding dimension is 192. At the same time, the sequence processed by the encoder is downsampled through the patch merging layer, and the embedding dimension is expanded to 256. Specifically, the resolution of the input sequence representation drops from to , and the number of channels changes from 192 to 256 through linear mapping.

[0051] The composition of the second stage is the same as that of the first stage. In the second stage, the height and width of the feature map are divided by 2 again, that is, 8x8, and then passed into the linear embedding layer. However, the number of heads of the second-stage encoder is 4, the number of layers is 2, and the embedding dimension is 256. The sequence processed by the encoder is downsampled through the patch merging layer, and the resolution drops from to , and at the same time, the embedding dimension is expanded to 384. Among them, the encoder in the second stage continues to enhance the global correlation between features, helping the model gradually focus on higher-level and more abstract feature representations.

[0052] In the third stage, the height and width of the feature map are divided by 2 again, that is, 4x4, and then passed into the linear embedding layer. The number of heads of the third-stage encoder is 6, the number of layers is 4, and the embedding dimension is 384. The sequence processed by the encoder is downsampled through the patch merging layer, and the resolution drops from to , while expanding the embedding dimension to 512, can improve the expression ability of features while downsampling.

[0053] In the fourth stage, the feature map is segmented according to the segmentation size of the feature map in the third stage, which is also 4x4, and then passed into the linear embedding layer. The number of heads of the encoder in the fourth stage is 8, the number of layers is 4, and the embedding dimension is 512. Among them, no downsampling is performed in the fourth stage, and the resolution of 4x4 remains unchanged. The main purpose is to further refine and fuse the high-level features extracted in the previous stages.

[0054] According to the above global average pooling function expression. If the final output shape is , then the global average pooling calculates to get the shape .

[0055] Furthermore, the fault diagnosis method also includes a method of using a hierarchical encoder to extract features from a sequence, including: Adding layer normalization to normalize the input sequence, calculating self-attention on the normalized sequence through the multi-head attention mechanism, adding the self-attention output to the original input sequence through a residual connection, and outputting the first feature with the same dimension as the input sequence. The function expression of the first feature output is:

[0056]

[0057] In the formula, represents the output first feature, X represents the input sequence, represents the normalized sequence, represents the output vector calculated by the attention function on the normalized sequence, represents the output vector calculated by the attention function, Q represents the query vector, K represents the key vector, V represents the value vector, represents a sparse vector with the same size as Q, D represents the key vector dimension, and softmax represents the classification function; Adding layer normalization to normalize the input first feature, applying a feed-forward layer to perform a non-linear transformation on the normalized first feature through a multi-layer perceptron, and adding the feed-forward layer output to the original input first feature through a residual connection, and outputting the second feature with the same dimension as the input first feature. The function expression of the second feature output is:

[0058] In the formula, represents the output second feature, represents the normalized first feature, It represents the output vector of the non - linear transformation of the first normalized feature by a multi - layer perceptron; The reference patch merging layer divides the second feature into multiple adjacent regions, splices adjacent sub - features within each region to obtain a composite feature, normalizes the composite feature along the channel dimension, and linearly projects and compresses the composite feature to the target number of channels to output the final feature map. The functional expression of the output of the final feature map is:

[0059] In the formula, It represents the output final feature map, - It represents the spliced sub - features, It represents the normalized composite feature vector, It represents the linear projection and compression of the composite feature to the target number of channels.

[0060] Through the implementation of the above - mentioned embodiments of the fault diagnosis method, first, the input sequence is normalized to stabilize the feature distribution and accelerate the training convergence. And the global dependencies within the sequence are captured through multi - head self - attention calculation to strengthen the associations at key positions. Then, the residual connection is used to retain the original input information, avoid gradient disappearance, enhance feature transmission, and ensure that low - level features are directly transmitted to the deep layers. On this basis, the feed - forward layer is applied to normalize the features enhanced by attention again to optimize subsequent processing, and a multi - layer perceptron is introduced for non - linear transformation to enhance the local feature expression ability. The attention and local features are fused through the residual connection to prevent network degradation. The output sequence is re - organized into a spatial feature map, and adjacent vectors are spliced within a local 2×2 region to halve the size of the feature map and reduce the computational complexity. And the composite feature is linearly projected and compressed to the target number of channels to double the number of channels and retain more semantic information. Thus, through the collaborative design of layer normalization, multi - head attention, residual connection, feed - forward layer, and patch merging, efficient global - local feature fusion is achieved, and multi - head attention and MLP are complementary, taking into account both long - range dependencies and local details. Normalization and residual connection ensure the effective propagation of gradients. Patch merging constructs hierarchical features to adapt to different task requirements.

[0061] As described above, combined with Figure 6 the unfolded diagram of the multi - head attention mechanism in

[0062] Preferably, for the bearing fault diagnosis model disclosed in the present invention, the tuning of hyperparameters is also crucial. The Adam optimizer and the CyclicLR learning rate scheduler can be used. The number of cycles is 50, and the batch size is 32. The Westinghouse bearing dataset is used for experiments, and a correct rate of 99.84% is obtained, demonstrating good diagnostic performance. At the same time, compared with other models, the correct rate of this model is also higher, as shown in the model comparison table: Model Comparison Table

[0063] In the above table, CoatNet (Convolution-Attention Network) is a hybrid architecture model that combines a convolutional neural network and a Transformer architecture; Alxnet is a deep convolutional neural network architecture model; Vgg16 (Visual Geometry Group 16) is a convolutional neural network architecture model; ViT (Vision Transformer) is a computer vision model based on the Transformer architecture. As shown in the model comparison table, the bearing fault diagnosis model of the embodiment of the present invention has a higher correct rate compared with other models, and at the same time, the number of parameters is much less than that of other traditional models.

[0064] The present invention also discloses a fault diagnosis system, including: A data preprocessing module, configured to obtain vibration data of the bearings of the target device, extract sample signals from the vibration data and divide them into a training set and a test set, and convert the sample signals of the training set and the test set into spectrograms respectively; A basic network construction module, configured to use two layers of lightweight convolutional layers and a pooling layer as the network basis, and combine a channel attention module to construct a feature enhancement module; A self-attention mechanism construction module, configured to add a patch partition layer after the feature enhancement module, and apply a multi-stage hierarchical multi-head attention mechanism after the patch partition layer. Hierarchical encoders with different numbers of layers, different numbers of heads, and different embedding dimensions are configured between each stage; A model establishment and training module, configured to build an improved vit architecture model with the feature enhancement module, the patch partition layer, and the hierarchical multi-head attention mechanism, input the spectrograms in the training set into the vit architecture model for model training, and obtain a bearing fault diagnosis model; A fault diagnosis module, configured to input the spectrograms in the test set into the trained bearing fault diagnosis model, and output the corresponding fault diagnosis results.

[0065] Based on the above fault diagnosis method, the present invention also discloses a computer-readable storage medium, and when the computer program is executed by a processor, the steps of the above fault diagnosis method are implemented.

[0066] Based on the above fault diagnosis method, the present invention also discloses a computer device, including a memory, a processor, and a computer program stored on the memory and executable on the processor. When the processor executes the computer program, the steps of the above fault diagnosis method are implemented.

[0067] The present invention is described in terms of flowcharts and / or block diagrams of methods, apparatuses (systems), and computer program products according to specific embodiments. It should be understood that each flow and / or block in the flowchart and / or block diagram, as well as the combination of flows and / or blocks in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to the processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing devices to generate a machine, such that the instructions executed by the processor of the computer or other programmable data processing devices generate a system for implementing the functions specified in Figure 1 one or more of the flows Figure 1 or one or more of the blocks.

[0068] These computer program instructions can also be stored in a computer-readable memory that can direct a computer or other programmable data processing device to work in a specific manner, such that the instructions stored in the computer-readable memory generate a manufactured article including an instruction system that implements the functions specified in Figure 1 one or more of the flows Figure 1 or one or more of the blocks.

[0069] These computer program instructions can also be loaded onto a computer or other programmable data processing device, such that a series of operation steps are executed on the computer or other programmable device to generate a computer-implemented process, and thus the instructions executed on the computer or other programmable device provide steps for implementing the functions specified in Figure 1 one or more of the flows Figure 1 or one or more of the blocks.

[0070] It should be understood that the above embodiments are only used to illustrate the technical solutions of the present invention and are not intended to limit them. For those skilled in the art, the technical solutions recorded in the above embodiments can be modified, or some of the technical features can be equivalently replaced; and all such modifications and replacements should fall within the protection scope of the appended claims of the present invention.

Claims

1. A fault diagnosis method based on an improved vit architecture with a hierarchical attention mechanism, characterized in that The described fault diagnosis method includes: Obtain the vibration data of the bearings of the target device, extract sample signals from the vibration data and divide them into a training set and a test set, and convert the sample signals of the training set and the test set into spectrograms respectively; Based on two layers of lightweight convolutional layers and one pooling layer as the network foundation, combine a channel attention module to construct a feature enhancement module; Add a patch partition layer after the feature enhancement module, and apply a multi-stage hierarchical multi-head attention mechanism after the patch partition layer. Hierarchical encoders with different numbers of layers, different numbers of heads, and different embedding dimensions are configured between each stage; Build an improved vit architecture model with the feature enhancement module, patch partition layer, and hierarchical multi-head attention mechanism. Input the spectrograms in the training set into the vit architecture model for model training to obtain a bearing fault diagnosis model; Input the spectrograms in the test set into the trained bearing fault diagnosis model, and output the corresponding fault diagnosis results.

2. The fault diagnosis method based on the improved vit architecture with a hierarchical attention mechanism according to claim 1, characterized in that The conversion of the sample signals of the training set and the test set into spectrograms respectively includes: Screen out the horizontal vibration signals from the vibration data, extract sample signals from the horizontal vibration signals with a preset step size and overlap rate, and randomly divide the extracted sample signals into a training set and a test set; Extract signal features from different frequencies and time scales for the divided sample signals through continuous wavelet transform, and convert them into spectrograms with preset pixel values. The function expression of the continuous wavelet transform is: In the formula, represents the mother wavelet function, s represents the scale factor, and τ represents the time shift factor, represents the time variable, represents the mother wavelet function after dilation and translation.

3. The fault diagnosis method based on the improved vit architecture with hierarchical attention mechanism according to claim 1, characterized in that, The described fault diagnosis method also includes a method of processing the input spectrogram using two layers of lightweight convolutional layers, including: Use the first convolutional layer to perform a convolution operation on the input spectrogram with a preset convolution kernel and step size. While keeping the size of the output feature map the same as that of the input spectrogram, expand the output channels to the number of the first channels; Add a batch normalization layer after the first convolutional layer to normalize the eigenvalue of each output channel, and introduce a Relu activation function after the batch normalization layer to suppress negative activation of each normalized eigenvalue, obtaining a first type of feature map after local feature extraction; Use the second convolutional layer to perform a convolution operation on the input first type of feature map with the same convolution kernel and step size. While keeping the size of the output feature map the same as that of the input feature map, expand the output channels to the number of the second channels, and output a second type of feature map after local feature extraction. The number of the second channels is greater than the number of the first channels; Add a batch normalization layer after the second convolutional layer to normalize the eigenvalue of each output channel, and introduce a Relu activation function after the batch normalization layer to suppress negative activation of each normalized eigenvalue, obtaining a second type of feature map after local feature extraction.

4. The fault diagnosis method based on the improved vit architecture with hierarchical attention mechanism according to claim 3, characterized in that The described fault diagnosis method also includes a method of performing convolution on a given position on the input spectrogram or feature map, including: According to the convolutional layer used, traverse each position of the convolution kernel, and add zero padding to the edge of the input spectrogram or feature map using padding; Slide the convolutional kernel over the input spectrogram or feature map to a given position, and calculate the output feature map at the given position. The functional expression for the output of the feature map at the given position is: Wherein, k represents the output channel index, i represents the index of the corresponding output feature map in the height direction, j represents the index of the corresponding output feature map in the width direction, K represents the convolution kernel size, u represents the position index of the convolution kernel in the height direction, and v represents the position index of the convolution kernel in the width direction. represents the input spectrogram or feature map, c represents the input channel number, and p represents the offset generated by using padding. represents a given position on the input spectrogram or feature map. the output feature map after convolution. represents the convolution kernel weight. represents the bias. represents the number of channels of the input spectrogram or feature map.

5. The fault diagnosis method based on the improved vit architecture with hierarchical attention mechanism according to claim 4, characterized in that, The fault diagnosis method further includes a method of using a pooling layer to process the input second type of feature map, including: Use a pooling layer with a pooling kernel of a preset size to perform pixel padding on the input second type of feature map; Slide a pooling window of a preset size over the padded second type of feature map at a preset stride, and extract the maximum eigenvalue within the pooling window after each slide, and output the pooled third type of feature map. The functional expression for the output of the pooled third type of feature map is: In the formula, represents the third type of feature map output after the pooling operation at a given position on the second type of feature map , represents the second type of feature map after padding, and m and n represent the relative offsets within the pooling window.​ 6. The fault diagnosis method based on the improved vit architecture with hierarchical attention mechanism according to claim 5, characterized in that The fault diagnosis method further includes a method of using a channel attention module to process the input third type of feature map, including: Extract the features of each channel of the third type of feature map, introduce an adaptive weight assignment layer on each channel, and the adaptive weight assignment layer dynamically adjusts the weights of each channel according to the feature responses of the channels; At the same time, use global average pooling and global max pooling to capture the statistical information of the features of each channel, send the obtained statistical information to a shared multi-layer perceptron, introduce a first fully connected layer to map the statistical information of global average pooling to the weight values of each channel, and introduce a second fully connected layer to map the statistical information of global max pooling to the weight values of each channel; Dynamically adjust the weights of the feature responses of each channel according to the mapped weight values, and perform weighted combination on the features of each channel through the mapped weight values, and output the channel attention map. The functional expression for the output of the channel attention map is: Wherein, Channel attention module function, represents the input data, represents the first fully connected layer, is the second fully connected layer, is function processing, represents the weight of the first fully connected layer, represents the result of global average pooling, represents the weight of the second fully connected layer, represents the result of global max pooling.

7. The fault diagnosis method based on the improved vit architecture with a hierarchical attention mechanism according to claim 6, characterized in that, The fault diagnosis method further includes a method of using a patch partition layer to process the input channel attention map, including: Set the partition window size of the patch partition layer, and slide the partition window size over the channel attention map to cut the channel attention map into multiple non-overlapping image patches along the spatial dimension, and output the partitioned fourth type of feature map.

8. The fault diagnosis method based on the improved vit architecture with hierarchical attention mechanism according to claim 7, characterized in that, The fault diagnosis method further includes a method of using a hierarchical multi-head attention mechanism to process the input fourth type of feature map, including: Add a linear embedding layer to perform linear projection on the input fourth type of feature map, map it to an embedding vector of a fixed dimension, then perform spatial compression and expand the spatial dimension, adjust the dimension order to obtain an embedding sequence, and use a hierarchical encoder with a preset number of layers, number of heads, and embedding dimension to initially extract features from the embedding sequence, and apply downsampling to increase the feature dimension, and then output the first-stage feature map; Add a linear embedding layer to perform spatial compression on the input first-stage feature map and flatten it into a sequence, use a hierarchical encoder with a preset number of layers, number of heads, and embedding dimension to further extract features from the flattened sequence. The number of heads and embedding dimension of the second-stage hierarchical encoder are greater than those of the first-stage hierarchical encoder, and apply downsampling to increase the feature dimension, and then output the second-stage feature map; Add a linear embedding layer to perform spatial compression on the input two-stage feature map and flatten it into a sequence, and use a hierarchical encoder with preset number of layers, number of heads, and embedding dimension to deeply extract features from the flattened sequence. The number of layers, number of heads, and embedding dimension of the three-stage hierarchical encoder are greater than those of the two-stage hierarchical encoder. After applying downsampling to increase the feature dimension, output the three-stage feature map; Add a linear embedding layer to perform spatial compression on the input three-stage feature map and flatten it into a sequence, and use a hierarchical encoder with preset number of layers, number of heads, and embedding dimension to establish global relationships for the input sequence, and output the high-level feature sequence after fusing the first three stages. The number of heads and embedding dimension of the four-stage hierarchical encoder are greater than those of the three-stage hierarchical encoder; Perform normalization on all the output high-level feature sequences, and perform global average pooling on the normalized high-level feature sequences along the sequence length dimension. The function expression of the global average pooling is: In the formula, represents the result after global pooling, represents the length after the feature map is flattened, represents the i-th eigenvalue of the input; Use a fully connected layer to map the globally average-pooled high-level feature sequences to a preset output dimension, and output the corresponding fault diagnosis results. The function expression of the mapping of the high-level feature sequences is: Wherein, The result of the mapped output after global average pooling, The weight matrix in the learning parameters, The bias in the learning parameters.

9. The fault diagnosis method based on the improved vit architecture with a hierarchical attention mechanism according to claim 8, characterized in that, The fault diagnosis method further includes a method for extracting features from a sequence using a hierarchical encoder, including: Add layer normalization to normalize the input sequence, calculate self-attention for the normalized sequence through a multi-head attention mechanism, add the self-attention output to the original input sequence through a residual connection, and output the first feature with the same dimension as the input sequence. The function expression of the output of the first feature is: In the formula, represents the first feature of the output, X represents the input sequence, represents the normalized sequence, represents the output vector calculated by the attention function for the normalized sequence, represents the output vector calculated by the attention function, Q represents the query vector, K represents the key vector, and V represents the value vector, represents a sparse vector of the same size as Q, D represents the key vector dimension, and softmax represents the classification function; Add layer normalization to normalize the input first feature, apply a feed-forward layer to perform a non-linear transformation on the normalized first feature through a multi-layer perceptron, and add the feed-forward layer output to the original input first feature through a residual connection, and output the second feature with the same dimension as the input first feature. The function expression of the output of the second feature is: In the formula, represents the second feature of the output, represents the first feature after normalization, represents the output vector of the non-linear transformation of the first feature after normalization by a multi-layer perceptron; Reference a patch merging layer to splice adjacent output vectors within a local area to obtain a composite feature vector, perform normalization on the composite feature vector along the channel dimension, and use a fully connected layer to compress the normalized composite feature vector to the target number of channels, and output the final feature map. The function expression of the output of the final feature map is: In the formula, represents the final feature map of the output, represents the output vector, represents the normalized composite feature vector, represents the compression of the composite feature vector to the target number of channels.

10. A fault diagnosis system, characterized in that, The fault diagnosis system includes: A data preprocessing module for obtaining the vibration data of the bearings of the target device, extracting sample signals from the vibration data and dividing them into a training set and a test set, and converting the sample signals of the training set and the test set into spectrograms respectively; A basic network construction module for using two layers of lightweight convolutional layers and a pooling layer as the network basis, and combining a channel attention module to construct a feature enhancement module; A self-attention mechanism construction module for adding a patch partitioning layer after the feature enhancement module, and applying a multi-stage hierarchical multi-head attention mechanism after the patch partitioning layer, and hierarchical encoders with different numbers of layers, different numbers of heads, and different embedding dimensions are configured between each stage; Model establishment and training module; used to build an improved vit architecture model with the feature enhancement module, patch partition layer, and hierarchical multi-head attention mechanism, input the spectrograms in the training set into the vit architecture model for model training, and obtain a bearing fault diagnosis model; Fault diagnosis module, used to input the spectrograms in the test set into the trained bearing fault diagnosis model and output the corresponding fault diagnosis results.

Images