Wavelet transform is widely used in vibration signal processing of rotating machinery due to its excellent time-frequency analysis ability for unsteady signals. The commonly used expression of wavelet transform is shown in Eq. 1. W T a , τ = 1 a ∫ − ∞ ∞ f t ψ t − τ a d t (1) where W T a , τ is the wavelet coefficient, f t is the input signal, ψ t is the wavelet basis function, a is the scale factor that controls the expansion and contraction of the wavelet basis function, and τ is the translation amount that controls the translation of the wavelet basis function. The scale corresponds to frequency, and the amount of translation corresponds to time. When the wavelet function is translated on the time axis at each scale, it is multiplied by the input signals respectively. So, the frequency components of the signals in each time period can be obtained.

The selection of wavelet basis function is an important step in using wavelet transform. Similarity coefficients (Mao et al., 2019; Liang et al., 2020) are applied to select the wavelet basis function for quantitative analysis. The expression of similarity coefficients is shown in Eq. 2. δ = ∑ i = 1 k α i m i 2 s i (2) where δ is the dimensionless similarity coefficient, k and s _i , m _i , α _i is the number of peaks, area of each peak, maximum of each peak, weighted coefficient of each peak after making absolute value for wavelet basis function.

Table 1 lists the similarity coefficients of several commonly used wavelets for fault diagnosis. When the similarity coefficient of a wavelet increases, the wavelet gets close to the original signals, which means the wavelet contains more fault information to facilitate diagnosis. Therefore, the cmor wavelet defined in Eq. 3 is used. The complex Morlet wavelet is a complex sinusoid modulated by a Gaussian envelope defined by: c m o r ( x ) = ( π F b ) − 1 2 ⁡ exp ( 2 i π F c x ) exp − x 2 F b (3) where F _b is the bandwidth of the wavelet, F _c is the central frequency of the wavelet function, and i is the imaginary unit. This paper takes F _c = F _b = 3.

Time-frequency images need to be preprocessed to improve the training performance before inputting into the network. For image datasets, the commonly used data enhancement methods include rotation, flipping, and random cropping. As the purpose of data enhancement, diverse training samples enable the network to extract key features from the samples undergoing various changes for improving the generalization ability of the network. However, since the vibration signal sequence is time-dependent, it can be regarded as a discrete time function. Flipping, rotation, and random cropping disrupt the relationship between time and frequency features, resulting in poor training performance. Therefore, the proposed image processing method first transforms the size of input images, and then normalizes them to improve the training speed.

Bilinear interpolation is used to convert the resolution of time-frequency images. Bilinear interpolation is the expansion of two-variable linear interpolation, which performs linear interpolation in two directions to obtain the value of the unknown function f (x, y) at the point P = (x, y). The transformation expression of bilinear interpolation is shown in Eq. 4. f ( x , y ) = ∑ i = 1 2 ∑ j = 1 2 f ( Q i j ) ( x i − x ) ( y j − y ) ( x 2 − x 1 ) ( y 2 − y 1 ) (4) where Q ₁₁ (x ₁, y ₁), Q ₂₁ (x ₂, y ₁), Q ₁₂ (x ₁, y ₂), and Q ₂₂ (x ₂, y ₂) are the four points closest to the target point f (x, y), and f (Q _ij ) represents the value of the point Q _ij , x _i,j , and y _ij represent the abscissa and ordinate of the point Q _ij , respectively.

Image normalization is used to scale the value of each image pixel to a small specific interval, remove the data unit limit, and convert it into a dimensionless pure value, which facilitates the comparison, and weighting of the indicators of different units or magnitudes. The normalization formula used in this paper is shown in Eq. 5. x ̄ = x − x min x max − x min (5) where the pixel value of a point is converted from x to x ̄ , x _max and x _min are the maximum and minimum values of the sample image pixel values, respectively. The values of all points can be converted to the interval of 0–1 by normalization, which improves the convergence speed and accuracy of the model.

According to the feature fusion method proposed in (Dong et al., 2020), the MSFF module is proposed to apply projection and back projection operators to the fault diagnosis of planetary gearboxes. As shown in Figure 1, the MSFF module fuses the features of all layers to supplement important time-frequency and fault information. Figure 3 shows the structure of the MSFF module. The MSFF module of the n-th network layer is defined as Eq. 6. j ̃ n = D n ( j n , { j ̃ 1 , j ̃ 2 , … , j ̃ n − 1 } ) (6) where j ⁿ is the latent feature of the n-th feature extraction network layer, j ̃ n is the enhanced features obtained through dense fusion, and { j ̃ 1 , j ̃ 2 , … , j ̃ n − 1 } are the enhanced fusion features from n MSFF modules before this layer in the network. This paper uses the enhanced features j ̃ n − t , t ∈ { 1,2 , … , n − 1 } , t times to gradually improve the enhanced features j ⁿ . The specific update and improvement process is shown as follows.

1) As shown in Eq. 7, the difference e t n between the enhanced feature j t n of the t-th iteration and the t-th enhanced feature j ̃ t is calculated.

Unlike the traditional back projection techniques, the sampling operators q t n and p t n in the network are unknown. The proposed method uses strided convolution (deconvolution layer) to learn the downsampling (upsampling) operator in an end-to-end manner. In order to avoid introducing too many parameters, this paper uses 2 as the stride size and stacks the convolutional and deconvolutional layers of n − t-th layer together to achieve downsampling and upsampling learning in q t n and p t n .

In summary, the multi-scale feature fusion (MSFF) algorithm used to enhance fault features is described in Algorithm 1 .

Algorithm 1

Fault feature enhancement algorithm based on multi-scale feature fusion

The MSFF module achieves feature fusion in the feature extraction process, but the issue of incomplete information still exists when the network uses the extracted features to classify faults. Therefore, the FoM module is designed before the classifier part, and its definition fom is shown in Eq. 9. As shown in Figure 4A, the FoM module is used to convert the output features of each network layer to a uniform size and concatenate on the channel dimensions to achieve the fusion of predicted features. f o m ( y 1 , y 2 … y L ) = c a t ( v 1 , v 2 … v L ) (9) v i = a m p W × H ( y i ) (10) where cat represents the concatenation of each input feature on the channel dimensions, and amp _W×H represents the adaptive maximum pooling operation with an output size of W × H. As shown in Figure 4B, the features of each scale are converted into feature vectors with the same resolution and different channel numbers to represent the diagnosis result obtained by each layer. y 1 , y 2 … y L are the output features of the first to L-th (L = 5 in this paper) network layers. The FoM module performs the second-time fusion of fault features to prevent the feature vectors used for classification from affecting the diagnosis result due to insufficient or missing information.

In the training process of deep learning networks, the hyperparameter settings have considerable impact on the training performance. In the experiments, the batch size of the network input (the number of samples input to the network each time) affects the testing accuracy of the model. If the batch size is too large, the model is difficult to fit or the fitting performance is poor. If the batch size is too small, the model is difficult to converge, and the accuracy rate oscillates unevenly. Considering the number of samples and the testing results, 16 is chosen.

The learning rate is another important hyperparameter. The learning rate represents the updated stride size of network parameters. If the learning rate is too small, it causes too low training efficiency and too much time is spent on training. If the learning rate is too large, it leads to local optimization, loss of oscillation, and model failure to converge. The dynamic adjustment mechanism of learning rate is adopted. The initial learning rate is 0.001, the learning rate transformation coefficient is 0.9. For instance, the learning rate is multiplied by 0.9 for every 10 training cycles. It ensures a fast training speed in the initial stage of training. When the convergence speed of the model slows down, the learning rate is reduced to gradually approach the optimal network parameter values.

All experiments are implemented on a computer with a RTX3090 GPU, 16 GB RAM, and an Intel i710700 CPU. The implementation, training and testing of the network model are executed on the Pytorch1.7.0 deep learning framework. Matlab 2018a is used to divide the original vibration signals into sample sequences and generate time-frequency image samples.

This paper uses the softmax cross-entropy loss function. The corresponding loss calculation formula for a single sample is shown in Eq. 11. l o ( x , c ) = − log exp ( x c ) ∑ j exp ( x j ) = − x c + log ∑ j exp ( x j ) (11) where c is the label (the category index), the vector x is the prediction result of each category (the network output), and x _c represents the c-th element of x. As the network trains, x _c approaches 1, so the loss approaches 0. The loss of each batch is shown in Eq. 12. l o s s = ∑ i = 1 N l o ( x i , C i ) N (12)

The vectors x _i and C _i are the network output and label category of each sample respectively, and N is the batch size. The loss of each batch is the average loss of each sample.

In order to speed up the convergence speed and avoid falling into the local optimization, the momentum-driven stochastic gradient descent (Momentum-SGD) algorithm is used to optimize the network model, and the momentum parameter is set to 0.9.

In order to verify the effectiveness of the proposed method, the proposed MSDFN is compared with DWWC + DRN (Zhao et al., 2017), WT-CNN (Liang et al., 2020), and three traditional image classification networks: ResNet18, DenseNet201, VGG11. There are the following explanations about the comparative experiment.

Since the innovations of the proposed method focus on the fault diagnosis network, the networks proposed by existing methods are used in the comparative experiments. In comparative experiments, the data preprocessing methods of the original literature are not included, but the same dataset of time-frequency images is used. Therefore, this paper selects the networks proposed in Zhao et al. (2017) (the original literature uses wavelet packet coefficient matrix) and Liang et al. (2020) (the original literature uses time-frequency images) for comparison.

In addition, several traditional image classification networks are applied to the performance comparison of network models. DBN (Wang et al., 2018) is the source of the dataset used in this article. The listed results are from the original literature. Since the original network structure is not convenient to process time-frequency images and its training method is quite different from others, this paper does not test DBN. All the related comparative results are from the original literature (Wang et al., 2018). Comparative experimental results. Table 7 shows the accuracy of fault diagnosis on each dataset. Compared with DBN and CNN using the traditional structures, the residual network and MSDFN have obvious advantages in accuracy. Compared with the residual network, the proposed network uses fewer network layers to achieve a certain improvement in fault diagnosis accuracy. The accuracy of each network on Dataset F, the University of Connecticut dataset proves the effectiveness of the proposed gearbox fault diagnosis method. In addition, all networks perform better on Dataset F, which also reflects the different complexity of the dataset distribution. Figure 7 shows the accuracies of several CNN-based networks obtained in 10 times repeated experiments. MSDFN has a stabler and higher accuracy rate on randomly allocated datasets, which means that the proposed network has stronger robustness. Figures 8A,B respectively compare the accuracy and loss curves of the Dataset E during the training process of each network. The convergence speed of MSDFN is close to ResNet, and its accuracy fluctuation is small.

In this paper, the feature fusion module first is removed and then ablation study is performed on the complete model to verify the impact of the proposed multi-scale dense fusion mechanism on diagnosis performance. The different degrees of ablation models are shown in Table 8. After removing all the remaining parts of the feature fusion structure, the diagnosis network shown in Figure 1 is used as backbone, which is mainly composed of a convolutional layer, a residual block, and a fully connected layer. The complete network D is MSDFN. The diagnostic accuracy of different degrees of ablation models on each dataset is shown in Table 9. It can be seen that the feature fusion module has a positive effect on improving the accuracy of fault diagnosis. The diagnostic accuracy of different degrees of ablation models does not fluctuate greatly under different working conditions, but they are still affected by the working conditions. Different models show different diagnostic performance on different datasets. As the load increases, the diagnostic accuracy of models B, C, and D is also improving, but model A does not have this rule, which reflects that the feature fusion module increases the model sensitivity to changes in load conditions.

Figure 9 shows the confusion matrices of the diagnosis results obtained by each ablated network model on the testing set of Dataset E containing 1,200 images. After adding the feature fusion structure to the network, the accuracy rate is significantly improved. Additionally, due to the dense feature fusion at each network layer, the MSFF module plays a more important role than the FoM module. Although the projection and back projection operations of each layer increase the amount of calculation, the fault diagnosis accuracy is effectively improved. Three images incorrectly diagnosed by the complete network D belong to different categories and are predicted to be three different categories, which indicates that MSDFN has no obvious deviation in the diagnosis ability of different fault types. However, the obvious deviation in the diagnosis ability exists in networks A, B, and C. The integration of the two feature fusion modules, MSFF and FoM, effectively improves the unstable recognition ability of the network for various fault types. 240 images are randomly selected from the testing set of Dataset E to compose a dataset. According to the network with various degrees of ablation, Figure 10 shows the distribution of output features after t-sne dimensionality reduction. The distribution of output features intuitively shows the separability of output features. The fault feature extraction performance of each ablated network can be evaluated.

The output features of the first two feature extraction network layers of each network model show two parts, a large part and a small part, because the separability of features is low at this time. Since testing is carried out on a mixed dataset, there are obvious data distribution differences between the data of the mixed working conditions. This mainly reflects the load among L1 (0 N.m), L2 (1.4 N.m), L3 (2.8 N.m), and L4 (25.2 N.m), so the two parts in the scatter diagram show an approximate ratio of 3:1. The performance of the two feature fusion modules on fault diagnosis can be analyzed by comparing the feature distribution of networks A, B, C, and D. As the network depth increases, the differences caused by working conditions gradually decrease, so the features of the same type of faults gradually become concentrated. Different ablated networks show different recognition capabilities. Compare with the third-layer output of different networks, the introduction of the MSFF module enhances the ability to distinguish features. According to the output comparison of each layer of each model, the network D achieves the best performance.

The fault recognition capability of each network can be determined by observing the output of FoM module in each model. Compared with other models, the output features of FoM module in model A (backbone network) have many errors, and the features of each fault category are relatively close. As the complete network, the output features of FoM module in model D show high separability, and 100% accuracy is achieved on the set of 240 random samples.