An improved semi-supervised prototype network for few-shot fault diagnosis

2.1. Prototypical networks

Prototypical Networks (ProtoNet) is a typical metric-based meta-learning model. This network maps the data into different classes by the feature function, and then learns the similarity between each sample by the distance metric [27], as illustrated in Fig. 1. Specifically, since ProtoNet operates on a meta-learning framework, it adheres to a scenario training strategy. Firstly, the dataset is divided into multiple tasks using the “N-way K-shot” approach. This means that each task comprises N classes, with each class containing K samples. Moreover, each task consists of a support set S (training set) and a query set Q (test set). The task is divided into two parts during the training process: the meta-training set and the meta-test set. Simultaneously, the types of data in the two are different [28].

Next, the training process can be more accurately described as that the prototype network mainly learns a nonlinear feature mapping function $f_e(\bullet )$. It is parameterized as a neural network to map the prototype into a feature space. In this space, prototypes from the same class are closer, while prototypes from different classes are far away from each other. All parameters of prototype networks depend on the feature mapping function. To calculate the prototype $P_d$ of each class $d$, the average value of the mapped feature vector represents the prototype. The calculation is as follows:

where $P_d$ is the prototype of class $d$, $\left|\sum S_i\right|$ represents the total size of samples in class $d$, $x_i^s$ is the sample belonging to class $d$ and $y_i^s$ is the label corresponding to $x_i^s$.

Finally, the parameter e of the feature function is trained and updated by the query set samples under each meta-task. The classification of the sample is done by measuring the distance between the sample and all prototypes. The prediction probability of the sample belonging to the query set of class d is calculated as follows:

For the prototype of all query sets, the loss function is defined to update the average negative logarithm of all query set samples in a given training set:

where $T_q$ is the total number of samples in the query set $Q$. The training process updates the model parameters by minimizing $L_{\theta }$.

Fig. 1Schematic representation of a metric-based meta-learning model

2.2. Prototype purification

Due to the small amount of sample data and the limited number of labeled samples, the initial prototype cannot correctly represent the cluster centroid. To get a better prototype, the unlabeled data is used to adjust the initial prototype. It can effectively solve the problem of position deviation caused by the limited labeled support data. The prototype can be guided slightly by the unlabeled samples, thereby better aggregating the information into the initial prototype. The prototype purification process is shown in Fig. 2.

First, the unlabeled data is assigned probability labels. And then, the probabilities of unlabeled data are aggregated into the existing prototypes to obtain more information. The initial prototype $P_d\in R^Z$ is constructed for each class with the labeled support set as shown in Eq. (1). Therefore, the probability of labeling the sample for each class can be given as follows:

The $S_l$ is a dataset that supports centralized labeling. The unlabeled data probability is shown in Eq. (5):

Fig. 2Schematic diagram of prototype purification

The standard Euclidean distance function ${sE}_{dist}$ is defined as:

where $o\in R^Z$;$\mathrm{}{s}_j$ is the standard deviation corresponding to the prototype. For the refined prototype, two probability weight coefficients are defined as:

Whereupon the original prototype is assigned by weight coefficient corresponding to each category:

The $s_l$ and $s_u$ are the number of labeled and unlabeled samples in the support set, respectively.

At last, the process of refining the prototype is completed by repeatedly executing the Eq. (9) and continuously assigning the value of ${\tilde{p}}_d$ to $P_d$.

3.1. Learning style

The proposed method follows the way of semi-supervised few-shot learning in this paper. Foremost, the support set $S$ is divided into two parts: labeled sample $S_l$ and unlabeled sample $S_u$:

After that, the labeled samples $S_l$ are then randomly selected for “N-way K-shot” classification. The labeled samples after classification are input into the model to obtain the initial prototypes of n categories. Moreover, $s_u$ samples in the unlabeled support set $S_u$ are input into the convolution operation blocks (COBs) module for mapping into the feature space. The unlabeled samples are used to adjust the position of initial prototype through multiple iterations until full convergence. In addition, part of the remaining samples are extracted from the $N$ classes as a query set $Q$, and then classify the samples in the query set $q_n$ by prototype purification. The prediction formula for each sample in the query set is given by:

where, $p_d\left(y=d\right)$ indicates the probability of a sample in the query set belonging to class $d$. Eventually, the model parameters are updated by calculating the loss of the model, and then back-propagated.

3.2. The diagnosis process of the proposed method

The general framework of the proposed method in this paper is shown in Fig. 3. The training process of method can be divided into four steps.

Step 1. Sample division and data input: The vibration acceleration signals of the bearing under different health conditions are collected. The meta-training set and the meta-test set contain multiple sets of 2048 data points generated by the time series signal, i.e., $x_i\in R^{1\times 2048}$. According to the learning method proposed in Section 3.1, the sample can be represented as $X={\left\{x_i\right\}}_{i=1}^{s_l+s_u+q_n}\in R^{\left(s_l+s_u+q_n\right)\times 1\times 2048}$. Where $s_l$ is a labeled sample in the support set $S_l$, $s_u$ is an unlabeled sample in the support set $S_u$, and $q_n$ is the sample in the query set.

Step 2. COBs block: The data is mapped to a feature space and then merge all the channel features to obtain a mapped feature vector of $f_e\left(x_i^s\right)$. Therefore, A convolution operation block (COBs) consisting of multiple cobs of the same type is constructed, which include convolution block, batch normalization block, activation function ReLU and pooling layer. Consequently, a COBs block is designed to facilitate feature extraction. The module is capable of mapping input data to a feature space. The details about the convolutional operation block can be seen in Table 1. The flow chart of the structure is shown in the upper right block diagram in Fig. 3.

Specifically, a convolution operation $F_{conv}$ is performed on the input data and the convolution kernel, thereby extracting the local input feature by sliding. The acceleration of the entire training process is achieved through the use of a batch normalization module. This approach also reduces the shift or disappearance of internal covariance, thus avoiding gradient explosion [29]. The input of $k$th BN layer can be set $x^k=\left\{x^{k\left(1\right)},x^{k\left(2\right)},\dots ,x^{k\left(M\right)}\right\}$, and the minibatch size is $M$, so $x^{k\left(m\right)}=\left\{x_1^{k\left(m\right)},x_2^{k\left(m\right)},\dots ,x_N^{k\left(m\right)}\right\}$. The BN operation can be shown as follows:

where ${\stackrel{-}{x}}_n^{k\left(m\right)}$ is BN layer output of the convolution result of the nth layer. ${\mu }_n$ and ${\sigma }_n^2$ are the mean and variance of $x_n^k$ respectively. Usually, $\epsilon$ is a very small constant to prevent invalid calculations when the variance is zero, and the value is usually $\epsilon =\mathrm{}$1×10^-5. Here, the parameters ${\gamma }_n^k$ and ${\beta }_n^k$ are respectively the scale parameter and the offset parameter to be learned. They are initially set to 1 and 0, respectively.

Fig. 3ISSPN overall structure diagram

To make the features learned by the convolution layer easier to distinguish, the activation function becomes an indispensable part after convolution operation. The ReLU function is a frequently used to activation function in deep neural networks. It can enhance the sparsity of network. Its left saturation function alleviates the gradient vanishing problem of the neural network to a certain extent and accelerates the convergence speed of gradient descent. The ReLU is actually a ramp function, which is defined as follows:

where ${\stackrel{-}{x}}_n^k$ is the output of upper layer (BN layer), and ${\stackrel{̿}{x}}_n^k$ is the activation amount of ${\stackrel{-}{x}}_n^k$ after activation by the ReLU function.

Further, the pooling layer is used to select features and reduce the dimension of features. The overfitting phenomenon can be effectively avoided by reducing the number of parameters. The maximum pooling $F_{MaxP}$ operation can better complete the current task, and it can complete the local maximum operation on the input feature map.

To sum up, set $W$ and $K$ to be the dimension and length of the channel, respectively. The output of $s$th convolution is $x_i^{\left(s\right)}\in R^{W\times K}$. Therefore, the above COBs can be described as:

Finally, all the channel features are combined. The input data is also mapped to the feature space. Therefore, it can be described as:

where $f_e\left(x_i\right)\in R^Z$, $Z=W\times K$ is a $Z$-dimensional eigenvector. The input data also completes the feature mapping process by the COBs block, so the result is $f_e\left(x_i^s\right)$.

Step 3 Prototype Purification: The process of refining the prototype is completed by repeatedly executing the Eq. (9) and continuously assigning the value of ${\tilde{p}}_d$ to $P_d$.To summarize, it can be generated the original prototype $P_d$ from the output obtained in the previous step. Meanwhile, the unlabeled data $s_u$ is assigned a predictable label ${\tilde{P}}_d$. Furthermore, it is aggregated into the original prototype to guide the original prototype with deviation for adjusting the position in multiple iterations until convergence, so as to purify the prototype.

Step 4. Update the model: The $q_n$ sample in the query set is purified by the prototype to obtain its prediction probability and realize classification all at once. Furthermore, the model parameters are updated by calculating the loss of $L_{\theta }$ and propagating backward.

Although the stochastic gradient descent (SGD) optimizer can quickly reduce the loss in the early stage of training, it often struggles to converge to an optimal minimum value for the model loss. The ISSPN is trained more efficiently by using adaptive moment estimation, which computes bias-corrected first and second moment estimates to counteract bias. Moreover, it automatically adjusts the learning rate scale for different layers, eliminating the need for manual selection as in the case of SGD. Hence, it is easier to use [30].

4.1. Description of the dataset

Due the fault signal characteristics of CWRU data set are very obvious. The validity of NSSPN method is further verified by collecting the bearing vibration signals of the comprehensive mechanical fault simulation test bench of the Spectral Mission (SQ) Company. Fig. 4 shows the integrated mechanical failure simulation test rig. The main part of the test bench is composed of driving motor, rotor system, motor control system, load and signal acquisition system. In the experiment, the vibration acceleration sensor is used to collect the bearing signal of the motor. The data acquisition instrument used is CoCo80, and the sampling frequency is 25.6 kHz. The experimental bearing is installed at the end of the drive motor, then a single point of failure is manually set on the experimental bearing. Specific information on the SQV bearing dataset is shown in Table 3. In addition, the experimental rotor system has four rotation frequencies (9 Hz, 19 Hz, 29 Hz, and 39 Hz). Fig. 5 shows a Time-Domain signal of the operating condition of the bearing at a rotation frequency of 39 Hz.

Fig. 4SQV bearing dataset test bench

4.2. Experimental parameter setting

To further evaluate the performance of the proposed method, two experiments are conducted on the SQV bearing dataset. The first experiment is three fault identifications (NC, IF-3 and OF-3) at four different speeds. First of all, a training data set is constructed: 100 samples are generated for each of the three fault conditions at each speed. So a total of 100×4×3 samples are generated. One to four samples are selected for each fault at each speed, resulting in four to sixteen samples for each type of fault. The number of unlabeled samples for ISSPN follows the support set, i.e. {1, 5} for each class. Ultimately, a single experiment and five experiments are set up to demonstrate the performance of the method, where each experiment is conducted for 30 iterations.

The second experiment is the identification of seven faults at one speed (39 Hz). Again, a data set needs to be constructed: 200 samples are generated for the 7 bearing States at 39 Hz, so a total of 1400 samples are generated. 5, 10, 20 and 30 samples are selected for each fault type. The training parameters are consistent with the above settings.

Fig. 5Time-domain signal of seven kinds of bearing vibration signals: a) mild inner ring fault (IF-1), b) moderate inner ring fault (IF-2), c) severe inner ring fault (IF-3), d) mild outer ring fault (OF-1), e) moderate outer ring fault (OF-2), f) severe outer ring fault (OF-3), g) normal condition (NC)

4.3. Experimental results and analysis

Table 4 shows the results of the three types of fault identification accuracy, and Fig. 6. visually shows the identification accuracy of each method. From the above tables and figures, it can be seen that although the accuracy of CNN, DaNN and ProtoNet is not less than 85 %, the ISSPN method shows better performance. The fault recognition rate of ISSPN can still reach 98.64 % when there are only 4 samples in each class. To sum up, the ISSPN method can adapt to fault diagnosis at different speeds by comparing with the other three methods.

Besides, the accuracy results for the seven types of faults are shown in Table 5 and Fig. 7. The ISSPN method also shows excellent performance in the identification of seven types of faults by comparing with the other three methods. Similarly, the fault identification of ISSPN is still excellent when there are the fewest fault samples of each type. So it can be seen that the ISSPN method can complete the multi-class fault diagnosis excellently.

Fig. 6Accuracy of identification of three types of faults on SQV dataset: a) one experiment, b) five experiment

a)

b)

Fig. 7Accuracy of identification of 7 types of faults on SQV dataset: a) one experiment, b) five experiment

a)

b)

Taking the classification of seven types of faults as an example, Fig. 8 and Fig. 9 show the classification characterization results of ISSPN method in the form of confusion matrix and t-SNE visualization, respectively. It can be seen from the confusion matrix that the ISSPN method can accurately complete the classification of seven types of faults. The classification results of ISSPN method are also more obvious in the t-SNE visualization map. In summary, the ISSPN method has better feature extraction ability.

Fig. 8Confusion matrix of ISSPN seven-class classification results on SQV dataset

Fig. 9Visualization of t-SNE for the four methods on SQV dataset