Soft Shared Multi-Task Convolutional Network Based on the Attention Mechanism

Rolling bearings are widely used and related to production safety, so bearing health status management is very important. In the bearing health management, fault diagnosis is the core problem, and residual lifetime (RUL) prediction is also very important. However, in the traditional single-task learning (STL) method, it is necessary to build a model for each task in which the training time is long and the efficiency is low. In addition, the traditional hard shared multi-task learning (MTL) method cannot be characterized according to each task characteristic, resulting in low generalization. Aiming at the problems of low efficiency of single task and low generalization hardware-shared multi-task capability, this paper proposes a soft shared multi-task convolutional network (SMTCNN) method based on attention mechanism and dynamic task weight to achieve accurate bearing fault diagnosis and RUL prediction. A shared feature extraction module convolutional neural network is established for joint training to improve efficiency. Then, an adaptive feature extraction module based on attention mechanism is built for different tasks to enhance the feature extraction ability of the model. Meanwhile, task weights are set dynamically based on uncertainty to enhance the model's generalization ability. Finally, compared with the STL and traditional hard sharing MTL, the proposed method achieves state-of-the-art results in terms of the fault diagnosis accuracy with 98.68% and the RUL prediction RMSE with 0.22, which proves the effectiveness of the proposed model in prognostics and health management.