基于多尺度1D-CNN卷积神经网络的轴承故障诊断代码+数据集

from torch import nn import torch import os import numpy as np import pandas as pd import matplotlib.pyplot as plt import time #####----绘制t-sne图----# from sklearn import datasets from sklearn.manifold import TSNE #####----绘制混淆矩阵图----# from sklearn.metrics import confusion_matrix from matplotlib.pylab import style #--- 定义预处理函数：加载数据、创建数据集---# def data_acquision(file_path): import scipy.io as scio data = scio.loadmat(file_path) # 加载mat数据 data_key_list = list(data.keys()) # mat文件为字典类型，将key变为list类型 accl_key = data_key_list[3] # mat文件为字典类型，其加速度列在key_list的第4个 accl_data = data[accl_key].flatten() # 获取加速度信号,并展成1维数据 accl_data = (accl_data-np.mean(accl_data))/np.std(accl_data) #Z-score标准化数据集 return accl_data def plot_freqDomain(data, fs=20480, title=None, clear_boundary=False, img_save_path=None, *args, **kwargs, ): import matplotlib.pyplot as plt import numpy as np ''' data = 要处理的数据 img_savepath = 保存图片路径 sampling_rate = 采样率 ''' fft_data = np.array(data) sampling_rate = fs ####-----加窗函数---#### fft_size = len(fft_data) window = np.hanning(fft_size) # 调节窗口种类 fft_data = fft_data*window ####----进行fft---#### xf = np.fft.fft(fft_data) xfp = np.fft.fftfreq(len(fft_data), d=1 / sampling_rate) # fftfreq(window length,) xfp = xfp.flatten() # 2维数据变成1维数据，后面要用。 xf = abs(xf) xf = xf[0:int(len(xf)/2)] #取xfp前半个长度 return xf def plot_envelopeSpectrum(data, fs=12000, title=None, clear_boundary=False, img_save_path=None, x1=None, x2=None, *args, **kwargs, ): ''' encelope_data = 要处理的数据 fs = 采样率 (default: 12000Hz) img_save_path = 图片保存路径，这里没用到 :return ET： 包络信号 ''' from scipy import fftpack import numpy as np envelope_data = np.array(data) # 获取数据 fft_size = len(envelope_data) # step1: 做希尔伯特变换 hx = fftpack.hilbert(envelope_data) # hx.shape = （14684，1） # step2: 求希尔伯特变换的幅值大小 hy = np.sqrt((envelope_data ** 2 + hx ** 2).astype('float')) # hy = sqrt(x^2 + hx^2) ##--------去直流分量 hy = hy - np.mean(hy) ##--------加窗函数 window = np.hanning(fft_size) # 调节窗口种类 hy = window * hy ET = np.fft.fft(hy) # 对希尔伯特变换后的hy做fft变换获得幅值 ET = np.abs(ET) # 对幅值求绝对值（此时的绝对值很大） freq = np.fft.fftfreq(fft_size, d=1 / fs) # 获取fft频率，ftfreq(fft_size,) freq = freq.flatten() # 2维数据变成1维数据，后面要用。 ## -----Z-score归一化 ET = (ET - np.mean(ET)) / np.std(ET) ET = ET[0:int(len(ET) / 2)] # 取xfp前半个长度 return ET def data_spilt(data, fr=None, num_2_generate=20, fs=None): ''' Desription: 将数据分割成n个小块。输入数据data采样点数是400000，分成100个子样本数据，每个子样本数据就是4000个数据点 param: data: 要输入的数据 fr: 转速(r/min) num_2_generate: 要生成的子样本数量 fs: 采样率 :return spilt_datalist 分割好的数据，类型为2维list（num_2_generate*each_fft_length ''' data = list(data) total_length = len(data) start_num = 0 # 子样本起始值 each_subdata_length = 2048 # 分割的子样本长度 end_num = each_subdata_length # 子样本终止值 step_length = int((total_length - each_subdata_length) / (num_2_generate - 1)) # step_length: 向前移动长度 i = 0 spilt_datalist = [] while i < num_2_generate: each_data = data[start_num: end_num] each_data = plot_freqDomain(fs=fs, data=each_data) # 输入fft数据 # each_data = each_data[-1:-1025:-1] #输入Time domain 数据 # each_data = plot_envelopeSpectrum(data=each_data, fs=20480) #输入ES domain 数据 each_data = (each_data-np.mean(each_data))/(np.std(each_data)) # 做Z-score归一化 spilt_datalist.append(each_data) start_num = 0 + i * step_length; end_num = each_subdata_length + i * step_length i = i + 1 return spilt_datalist #-----创建Dataset----# # base_dir = r'E:/研究生/pytorch/CWRU-MSCNN (2)' train_data_base_dir = os.path.join(base_dir,'dataset','train_data') # 这里需要改成自己的路径 fault_type_list = os.listdir(train_data_base_dir) file_path_list = list() train_data = [] train_label = [] data_base_dir = r'E:/研究生/pytorch/CWRU-MSCNN/dataset/train_data' fault_type_list = os.listdir(data_base_dir) file_path_list = list() train_data = [] train_label = [] for fault_type in fault_type_list: file_name_list = os.listdir(os.path.join(data_base_dir, fault_type)) for file_name in file_name_list: print(file_name) num_2_generate = 60 fs=12000 if 'Normal' in file_name: num_2_generate = num_2_generate*6 # Normal样本数量比较少，进行扩充 if '48k' in file_name: fs = 48000 file_path = os.path.join(data_base_dir, fault_type, file_name) fault_type2fault_label = {'BF':'0', 'OF':'1', 'IF':'2', 'Normal':'3'} # 获取数据 data = data_acquision(file_path) # 获取子样本数据 sub_data_list = data_spilt(data=data, num_2_generate=num_2_generate,fs=fs) train_label.extend(list(fault_type2fault_label[fault_type]*len(sub_data_list))) train_data.extend(sub_data_list) test_data_base_dir = os.path.join(base_dir,'dataset','test_data') # 这里需要改成自己的路径 fault_type_list = os.listdir(test_data_base_dir) file_path_list = list() test_data = [] test_label = [] for fault_type in fault_type_list: file_name_list = os.listdir(os.path.join(test_data_base_dir, fault_type)) for file_name in file_name_list: print(file_name) num_2_generate = 60 fs=12000 if 'Normal' in file_name: num_2_generate = num_2_generate*6 fs = 48000 file_path = os.path.join(test_data_base_dir, fault_type, file_name) fault_type2fault_label = {'BF':'0', 'OF':'1', 'IF':'2', 'Normal':'3'} # 获取数据 data = data_acquision(file_path) # 获取子样本数据 sub_data_list = data_spilt(data=data, num_2_generate=num_2_generate, fs=fs) test_label.extend(list(fault_type2fault_label[fault_type]*len(sub_data_list))) test_data.extend(sub_data_list) #----创建dataloader---# from torch.utils.data import DataLoader from torch.utils.data import TensorDataset train_data = np.array(train_data) train_label = np.array((train_label), dtype=int) train_x = torch.from_numpy(train_data).type(torch.float32) ##torch.Size([2880, 1024]) train_y = torch.from_numpy(train_label).type(torch.LongTensor) ##torch.Size([2880]) train_x = torch.unsqueeze(train_x, dim=1) ##扩展成3维 torch.Size([2880, 1, 1024]) test_data = np.array(test_data) #测试数据 test_label = np.array((test_label), dtype=int) #测试标签array([0, 0, 0, ..., 1, 1, 1]) test_x = torch.from_numpy(test_data).type(torch.float32) ##torch.Size([720, 1024]) test_y = torch.from_numpy(test_label).type(torch.LongTensor) ##torch.Size([720]) test_x = torch.unsqueeze(test_x, dim=1) ##torch.Size([720, 1, 1024]) #创建dataset from sklearn.model_selection import train_test_split train_ds = TensorDataset(train_x, train_y) train_ds,valid_ds = train_test_split(train_ds, test_size = 0.2,random_state = 42) #将训练集分隔维训练集和测试集，分割比例5:1 test_ds = TensorDataset(test_x, test_y) #创建dataloader batch_size = 64 train_dl = DataLoader(dataset = train_ds, batch_size = batch_size,