Python环境下基于原型网络的滚动轴承故障诊断方法

近年来，基于深度学习的故障诊断方法成为研究热点，它依靠海量的数据完成故障诊断模型的训练。然而，在实际的工业过程中，往往无法提供充足的故障数据样本给深度学习模型，在一定程度上限制了其性能。以旋转机械为例，其运行需要保持高稳定性与可靠性，其从正常运行状态过渡到故障状态需要经历漫长的时间，且其故障状态持续时间较短，因此难以采集到足够的故障数据。此外，旋转机械设备的工况易发生变化，采集足够的数据并对其进行标签化是比较困难的，且较为耗费人力。因此，实际工业中采集到的多数是正常数据，故障状态下的数据较少，使得模型无法正确学习到从振动信号到故障类别的映射关系，导致模型欠拟合，泛化能力差。

为解决数据容量限制深度学习应用的问题，小样本学习应运而生。小样本学习又称为少样本学习，它的概念最早从计算机视觉领域兴起。近年来小样本学习得到了广泛关注，并被应用于图像处理和故障诊断等领域。小样本学习的目标是在每一类数据较少的情况下，从匮乏的数据中学习到解决问题的办法。目前，小样本学习的方法主要分为三类：基于优化的方法，基于数据增强的方法和基于度量学习的方法。

基于优化的方法是小样本问题的经典解决方法，该方法通过大规模的源数据集学习特征以得到泛化能力很强的模型，再通过目标数据集对模型进行微调。常用的基于优化的方法包含模型无关的元学习算法MAML和Reptile等。然而，在实际场景中源数据集和目标数据集分布不一定类似，基于优化的方法会导致模型在目标数据集上产生过拟合的问题。此外，基于优化的方法在寻找最优参数的时候需要人工经验对梯度的组合方式进行选择，上述问题限制了基于优化的方法的应用场景。基于数据增强的方法的目的是对原始数据集进行扩充，使数据容量能够满足深度学习模型的需求，是解决小样本学习问题的直接方法。度量学习也称为相似度学习，通过将原始输入映射到嵌入空间得到嵌入特征，并使用某种距离度量函数计算特征之间的差异，从而完成分类任务。

1）基于数据增强的故障诊断算法小样本学习的根本问题在于数据量过少，导致样本的多样性降低，从而限制了深度学习模型的性能。而基于数据增强的方法是一种提高训练集数量和质量的技术，可以在一定程度上增加样本多样性。它可以利用原有数据中的信息，在嵌入空间中对数据进行增强以生成新的数据，该方法能够直接解决小样本学习中数据匮乏的问题。常用的数据增强方法包含数据扩充和特征增强。其中，前者主要是使用无参数的方法，包括对图像进行剪切、旋转等操作，以及对原始数据添加噪声等方式。后者则是训练出深度神经网络，学习原始数据与目标数据之间的映射关系，在特征空间中增加便于分类的特征。

2）基于度量学习的故障诊断算法

针对少量样本的情况，可采用基于数据增强的方法对故障样本进行扩充，再使用扩充后的数据集进行故障诊断。但是对于没有历史累计数据的旋转机械，其故障样本可能只有几十个，这时旋转机械处于极少量故障样本的情形，数据集为极小样本数据集。由于其他文献中少量样本和极少量样本无明确的界限与定义，因此在综合考虑滚动轴承数据的类别数和时频图差异性后，可以对少量样本和极少量样本给出一个简单的定义。其中，少量样本为每个类别的样本数大于等于50且小于等于100的情形，极少量样本为样本数小于50的情形。在极少量样本情形下，数据增强的方法受限，生成样本的多样性无法提升，从而影响故障诊断精度。而度量学习通过特征表达和相似性度量，减少了网络需要学习的参数，从而减弱样本稀缺对模型造成的影响。因此，在极少量样本的情况下使用基于度量学习的方法进行故障诊断是一个比较好的方案。

基于度量学习的方法主要研究如何在特定任务上学习一个距离函数，该函数可以对样本对之间的相似性进行度量，从而使基于近邻思想的算法得到较好的模型参数。度量学习中的深度度量学习可以学习原始数据与嵌入空间的映射关系，使得同类别的样本对距离较近，而不同类别的样本对距离较远。通过使用距离函数度量待分类样本与其他样本的距离，便可完成分类任务。其中，基于度量学习的方法中具有代表性的网络为匹配网络、孪生网络、原型网络和关系网络等。由于其出色的分类性能，近年来基于原型网络的故障诊断方法得到了较多关注，逐渐成为研究的热点。

原型网络是一种基于度量学习的小样本学习方法。该方法旨在学习各类样本在一个度量空间的原型表示，通过比较查询样本与各类原型之间的距离，将查询样本归入距离最近的类别，从而达到分类的目的。

鉴于此，提出一种基于原型网络的滚动轴承故障诊断方法，运行环境为Python，可能采用的数据集如下：

数据集A-D由西储大学CWRU轴承数据集而来，数据集E-F由KAt数据集而来，主代码如下：

import numpy as np
import scipy.io as sio 
import os
import Prototypical_Network as PN
import random


def get_data(datapath, is_normalization=False, is_fft=False):
    data = sio.loadmat(datapath)
    
    try:
        trainX = data['trainX'].astype(np.float32) 
        testX = data['testX'].astype(np.float32) 
        trainY = data['trainY']
        testY = data['testY']   
        if trainY.shape[0]>1:
            trainY = np.argmax(trainY, axis=1)
            testY = np.argmax(testY, axis=1)
        else:
            trainY = trainY-1
            testY = testY-1
    except:
        X = data['X'].astype(np.float32) 
        Y = np.argmax(data['Y'], axis=1)
        trainX, testX, train_indices, test_indices = utils.split_train_test(X, test_ratio=0.5)
        trainY, testY = Y[train_indices], Y[test_indices]
        
    if is_fft:
        trainX = np.absolute(np.fft.fft(trainX, axis=1)).astype(np.float32) 
        testX = np.absolute(np.fft.fft(testX, axis=1)).astype(np.float32) 

    if is_normalization:
        trainX, mu, sigma = utils.zscore(trainX, dim=0)
        testX = utils.normalize(testX, mu, sigma)
    return trainX, trainY, testX, testY



def get_sample_data(X, Y, num=10, classes=None):
    if classes is None: classes = len(np.unique(Y))
    
    Xs, Ys = [], [],
    for i in range(classes):
        idx = np.where(Y==i)
        idx = random.sample(list(idx[0]), num)
        
        Xs.extend(X[idx])    
        Ys.extend(Y[idx])    
    return np.array(Xs), np.array(Ys)

# ---------------------------------------------
#  few-shot learning with prototypical network
# ---------------------------------------------
datasets = ['dataA', 'dataB', 'dataC', 'dataD', 'dataE', 'dataF']
datapath = 'datasets/'
for i in [0]:
    utils.setup_seed(seed=0)
    data = datasets[i]
    save_result = f'./results/{data}'
    
    #utils.setup_seed(seed=0)
    trainX, trainY, testX, testY = get_data(datapath+data, is_fft=True)
    trainX, trainY = get_sample_data(trainX, trainY, num=10) # randomly select a small number of training data

    
    ### Define Nets
    encoder = PN.encoder()
    
    ### Train Nets
    model = PN.Protonet(trainX, trainY, n_epoch=200, is_normalization=True)
    model = model.train(encoder)
         
    ### test
    out = model.test(testX)
    
    result = {}
    result['distance'] = out['distance']
    result['prediction'] = out['prediction']
    result['embedding'] = out['embedding']
    result['loss'] = model.loss
    result['prototypes'] = model.prototypes.detach().cpu().numpy()
    result['true_label'] = testY
    
        
    if not os.path.exists(save_result): os.makedirs(save_result)
    sio.savemat(f'{save_result}/result-{model.n_support}shot.mat', result)   

结果如下：

完整代码可由知乎学术咨询获得：

Python环境下基于原型网络的滚动轴承故障诊断方法

工学博士，担任《Mechanical System and Signal Processing》审稿专家，担任《中国电机工程学报》优秀审稿专家，《控制与决策》，《系统工程与电子技术》，《电力系统保护与控制》，《宇航学报》等EI期刊审稿专家。

擅长领域：现代信号处理，机器学习，深度学习，数字孪生，时间序列分析，设备缺陷检测、设备异常检测、设备智能故障诊断与健康管理PHM等。