A class-aware supervised contrastive learning framework for imbalanced fault diagnosis

Deep learning-based fault diagnosis models constructed from imbalanced datasets would meet severe performance degradation when the number of samples for fault classes is much smaller than the normal category. Recent feature-learning-based methods have shown that learning the discriminative feature representation helps construct the well-performing fault classifier in imbalanced fault diagnosis. However, these methods have limited discriminative feature extraction abilities when applying them to a more practical but not well-studied class-imbalance scenario, where only the normal condition has a large amount of data while sample sizes of all fault classes are small. To address this issue, a novel feature-learning-based method called Class-aware Supervised Contrastive Learning (CA-SupCon) is proposed. Supervised contrastive learning (SupCon) is adopted for the first time in imbalanced fault diagnosis to optimize the feature difference between any two classes by leveraging category information. Additionally, a class-aware sampler (CA) is designed to rebalance data distribution within each mini-batch during training, which improves the ability of SupCon to enlarge the feature distance between any two minority fault conditions. By effectively integrating SupCon and CA, the proposed CA-SupCon framework can obtain a more discriminative feature space with better intra-class compactness and inter-class separability, and achieves good performance under the above class-imbalance scenario. Extensive experiments on two open-source datasets demonstrate the effectiveness of the proposed method. Code is available at https://github.com/JiyangZhang-UESTC/CA-SupCon.