Prior knowledge embedding convolutional autoencoder: A single-source domain generalized fault diagnosis framework under small samples

The proposed transfer learning-based fault diagnosis models have achieved good results in multi-source domain generalization (MDG) tasks. However, research on single-source domain generalization (SDG) is relatively scarce, and the limited availability of small training samples is seldom considered. Therefore, to address the insufficient feature extraction capability and poor generalization performance of existing models on single-source domain small sample data, a novel single-source domain generalization fault diagnosis (SDGFD) framework, the prior knowledge embedded convolutional autoencoder (PKECA), is proposed. During the training phase, first, single-source domain data are used to construct prior features based on the time domain, frequency domain, and time-frequency domain. Second, a prior knowledge embedding structure based on the convolutional autoencoder is built, which compresses the prior knowledge and original vibration data into a high-dimensional space of consistent dimensions, embedding the prior knowledge into the features corresponding to the vibration data using a mean squared error loss function. Subsequently, the proposed centroid-based self-supervised learning (CBSSL) strategy further constrains high-dimensional features, improving the generalization ability. The designed sparse regularized activation (SRA) function significantly enhances the regularization effect on features. During the testing phase, it is only necessary to input the data from the unknown domain to identify the fault types. The experimental results show that the proposed method achieves superior performance in fault diagnosis tasks involving cross-speed, time-varying speed, and small sample data in SDGFD, demonstrating that PKECA has strong generalizability. The code can be found here: https://github.com/John-520/PKECA. (c) 2024 Elsevier Science. All rights reserved