CMFE-PVT: a lightweight fault diagnosis framework for rolling bearings using compact multi-scale feature extraction and a pruned-restructured vision transformer

Abstract The intelligent fault diagnosis method based on transformer and convolutional neural network (CNN) has achieved good global and local feature extraction results. However, the multi-head self-attention mechanism adopted by the transformer and the cross-channel convolution operation in CNN increases the complexity of the model, thereby increasing the demand for hardware resources, which to some extent, limits its broad applicability in industrial applications. Therefore, this paper proposes a lightweight fault diagnosis framework based on compact multi-scale feature extraction and pruned-restructured vision transformer (ViT) to address the above challenges. Firstly, a compact multi-scale feature extraction module is designed to efficiently capture complex features in rolling bearing vibration signals through parallel multi-scale convolution kernels, combined with channel reduction strategies to significantly reduce computational complexity while maintaining feature richness. Next, short-time Fourier transform and pseudo-color processing techniques are used to obtain time–frequency images. Then, a dual optimization of matrix sparsity and structural reorganization is implemented for Self-attention in ViT to ensure model performance and significantly reduce computational overhead. Finally, the time–frequency images are segmented and rearranged before being fed into the improved lightweight ViT for global feature extraction and fault recognition of rolling bearings. The experimental results show that the proposed fault diagnosis method has the advantages of lightweight (Params:4.27 K, floating point operations per seconds:0.1 M, multiplication and accumulation operations per seconds:51.07 K) and robustness compared to mainstream algorithms.