滚动轴承故障检测 基于双路神经网络滚动轴承故障诊断系统 系统融合原始振动信号和二维时频图像进行多输入(多通道)故障诊断。PyTorch框架实现模型,可单双路

📅 2026/8/13 9:04:55
滚动轴承故障检测 基于双路神经网络滚动轴承故障诊断系统 系统融合原始振动信号和二维时频图像进行多输入(多通道)故障诊断。PyTorch框架实现模型,可单双路
基于双路神经网络的滚动轴承故障诊断文章目录基于双路神经网络的滚动轴承故障诊断文字及内容代码仅供参考1. 数据准备2. 完整代码实现2.1 数据加载和预处理2.2 模型构建2.3 模型训练2.4 结果可视化运行步骤实际运行情况如图融合了原始振动信号 和 二维信号时频图像 的多输入(多通道)故障诊断方法单路和双路都可时频图像算法可选小波变换短时傅里叶变换马尔可夫变迁场格拉姆角场S变换递归图等预处理不提供数据集可选凯斯西储大学东南大学江南大学等等使用 Pytorch和Tensorflow框架构建一个基于双路神经网络的滚动轴承故障诊断系统。该系统将融合原始振动信号和二维时频图像进行多输入多通道故障诊断。使用PyTorch框架来实现这个模型单路和双路两种模式都可文字及内容代码仅供参考1. 数据准备假设我们有一个滚动轴承振动信号数据集其中包含正常和故障样本。数据集格式如下datasets/ └── case_western_reserve_university/ ├── train/ │ ├── normal/ │ ├── outer_race_fault/ │ ├── inner_race_fault/ │ └── ball_fault/ ├── val/ │ ├── normal/ │ ├── outer_race_fault/ │ ├── inner_race_fault/ │ └── ball_fault/ └── test/ ├── normal/ ├── outer_race_fault/ ├── inner_race_fault/ └── ball_fault/每个文件夹中包含多个.npy文件每个文件是一个形状为(timesteps, features)的数组。同时我们还需要生成相应的二维时频图像。2. 完整代码实现以下是完整的Python代码包括数据加载、预处理、模型构建、训练和评估。您可以直接运行此代码。2.1 数据加载和预处理[titleData Loading and Preprocessing]importosimportnumpyasnpfromsklearn.preprocessingimportStandardScalerfromkeras.utilsimportto_categoricalimporttorchfromtorch.utils.dataimportDataset,DataLoaderimporttorchvision.transformsastransformsfromscipy.signalimportstftimportmatplotlib.pyplotasplt# Load data from .npy filesdefload_data(data_dir,classes):X[]y[]forlabel,clsinenumerate(classes):class_diros.path.join(data_dir,cls)filesos.listdir(class_dir)forfileinfiles:filepathos.path.join(class_dir,file)datanp.load(filepath)X.append(data)y.append(label)Xnp.array(X)ynp.array(y)returnX,y# Preprocess time-domain signalsdefpreprocess_signals(X_train,X_val,X_test):scalerStandardScaler()n_timesteps,n_featuresX_train.shape[1],X_train.shape[2]X_trainscaler.fit_transform(X_train.reshape(-1,n_features)).reshape(-1,n_timesteps,n_features)X_valscaler.transform(X_val.reshape(-1,n_features)).reshape(-1,n_timesteps,n_features)X_testscaler.transform(X_test.reshape(-1,n_features)).reshape(-1,n_timesteps,n_features)returnX_train,X_val,X_test# Generate STFT images from time-domain signalsdefgenerate_stft_images(X):stft_images[]forsignalinX:f,t,Zxxstft(signal[:,0],fs1000,nperseg256,noverlap128)magnitude_spectrogramnp.abs(Zxx)normalized_image(magnitude_spectrogram-np.min(magnitude_spectrogram))/(np.max(magnitude_spectrogram)-np.min(magnitude_spectrogram))stft_images.append(normalized_image)returnnp.array(stft_images)# Custom dataset classclassBearingDataset(Dataset):def__init__(self,X_signal,X_image,y,transformNone):self.X_signalX_signal self.X_imageX_image self.yy self.transformtransformdef__len__(self):returnlen(self.y)def__getitem__(self,idx):signalself.X_signal[idx]imageself.X_image[idx]labelself.y[idx]ifself.transform:imageself.transform(image)returnsignal,image,label# Load and preprocess classification dataclasses[normal,outer_race_fault,inner_race_fault,ball_fault]data_dirdatasets/case_western_reserve_universityX_train_class,y_train_classload_data(os.path.join(data_dir,train),classes)X_val_class,y_val_classload_data(os.path.join(data_dir,val),classes)X_test_class,y_test_classload_data(os.path.join(data_dir,test),classes)y_train_classto_categorical(y_train_class,num_classeslen(classes))y_val_classto_categorical(y_val_class,num_classeslen(classes))y_test_classto_categorical(y_test_class,num_classeslen(classes))X_train_class,X_val_class,X_test_classpreprocess_signals(X_train_class,X_val_class,X_test_class)# Generate STFT imagesX_train_imagegenerate_stft_images(X_train_class)X_val_imagegenerate_stft_images(X_val_class)X_test_imagegenerate_stft_images(X_test_class)# Create datasets and dataloaderstransformtransforms.Compose([transforms.ToTensor(),])train_datasetBearingDataset(X_train_class,X_train_image,y_train_class,transformtransform)val_datasetBearingDataset(X_val_class,X_val_image,y_val_class,transformtransform)test_datasetBearingDataset(X_test_class,X_test_image,y_test_class,transformtransform)train_loaderDataLoader(train_dataset,batch_size32,shuffleTrue)val_loaderDataLoader(val_dataset,batch_size32,shuffleFalse)test_loaderDataLoader(test_dataset,batch_size32,shuffleFalse)2.2 模型构建[titleModel Construction]importtorch.nnasnnimporttorch.nn.functionalasF# Single-channel CNN modelclassSingleChannelCNN(nn.Module):def__init__(self,input_shape,num_classes):super(SingleChannelCNN,self).__init__()self.conv1nn.Conv1d(in_channelsinput_shape[1],out_channels64,kernel_size5,padding2)self.poolnn.MaxPool1d(kernel_size2)self.dropoutnn.Dropout(p0.2)self.fc1nn.Linear(64*(input_shape[0]//2),128)self.fc2nn.Linear(128,num_classes)defforward(self,x):xself.pool(F.relu(self.conv1(x)))xself.dropout(x)xx.view(x.size(0),-1)xF.relu(self.fc1(x))xself.dropout(x)xself.fc2(x)returnx# Two-channel CNN BiLSTM Attention modelclassDualChannelCNNBiLSTMAtt(nn.Module):def__init__(self,input_shape_signal,input_shape_image,num_classes):super(DualChannelCNNBiLSTMAtt,self).__init__()# Signal branchself.signal_conv1nn.Conv1d(in_channelsinput_shape_signal[1],out_channels64,kernel_size5,padding2)self.signal_poolnn.MaxPool1d(kernel_size2)self.signal_dropoutnn.Dropout(p0.2)self.signal_lstmnn.LSTM(input_size64,hidden_size64,bidirectionalTrue,batch_firstTrue)self.signal_attentionself.attention_3d_block# Image branchself.image_conv1nn.Conv2d(in_channels1,out_channels64,kernel_size(5,5),padding2)self.image_poolnn.MaxPool2d(kernel_size(2,2))self.image_dropoutnn.Dropout(p0.2)self.image_flattennn.Flatten()self.image_fc1nn.Linear(64*(input_shape_image[0]//2)*(input_shape_image[1]//2),128)# Combined branchself.fc_combined1nn.Linear(128*2,256)self.fc_combined2nn.Linear(256,num_classes)defattention_3d_block(self,inputs):# inputs.shape (batch_size, seq_len, feature_dim)single_attention_vectorFalseattention_weightsF.tanh(self.dense(inputs))attention_weightsattention_weights.permute(0,2,1)attention_weightsF.softmax(attention_weights,dim-1)context_vectortorch.bmm(attention_weights,inputs)ifsingle_attention_vector:context_vectorcontext_vector.squeeze(1)else:context_vectorcontext_vector.permute(0,2,1)returncontext_vectordefdense(self,inputs):returnnn.Linear(inputs.shape[-1],1)(inputs)defforward(self,signal,image):# Signal branchsignalself.signal_pool(F.relu(self.signal_conv1(signal)))signalself.signal_dropout(signal)signal,_self.signal_lstm(signal)signal_attself.signal_attention(signal)signal_flatsignal_att.view(signal_att.size(0),-1)# Image branchimageself.image_pool(F.relu(self.image_conv1(image.unsqueeze(1))))imageself.image_dropout(image)image_flatself.image_flatten(image)image_outF.relu(self.image_fc1(image_flat))# Combined branchcombinedtorch.cat((signal_flat,image_out),dim1)combinedF.relu(self.fc_combined1(combined))combinedself.fc_combined2(combined)returncombined# Define modelsinput_shape_signalX_train_class.shape[1:]input_shape_imageX_train_image.shape[1:]single_channel_cnnSingleChannelCNN(input_shape_signal,num_classes)dual_channel_cnn_bilstm_attDualChannelCNNBiLSTMAtt(input_shape_signal,input_shape_image,num_classes)print(single_channel_cnn)print(dual_channel_cnn_bilstm_att)2.3 模型训练[titleModel Training]importtorch.optimasoptim# Function to train a modeldeftrain_model(model,train_loader,val_loader,criterion,optimizer,num_epochs50):devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)model.to(device)best_acc0.0best_model_wtscopy.deepcopy(model.state_dict())forepochinrange(num_epochs):print(fEpoch{epoch}/{num_epochs-1})print(-*10)# Each epoch has a training and validation phaseforphasein[train,val]:ifphasetrain:model.train()# Set model to training modeelse:model.eval()# Set model to evaluate moderunning_loss0.0running_corrects0# Iterate over data.forsignals,images,labelsintrain_loaderifphasetrainelseval_loader:signalssignals.float().to(device)imagesimages.float().to(device)labelstorch.tensor(labels.argmax(axis1)).long().to(device)# Zero the parameter gradientsoptimizer.zero_grad()# Forwardwithtorch.set_grad_enabled(phasetrain):outputsmodel(signals,images)_,predstorch.max(outputs,1)losscriterion(outputs,labels)# Backward optimize only if in training phaseifphasetrain:loss.backward()optimizer.step()# Statisticsrunning_lossloss.item()*signals.size(0)running_correctstorch.sum(predslabels.data)epoch_lossrunning_loss/len(train_loader.datasetifphasetrainelseval_loader.dataset)epoch_accrunning_corrects.double()/len(train_loader.datasetifphasetrainelseval_loader.dataset)print(f{phase}Loss:{epoch_loss:.4f}Acc:{epoch_acc:.4f})# Deep copy the modelifphasevalandepoch_accbest_acc:best_accepoch_acc best_model_wtscopy.deepcopy(model.state_dict())print()print(fBest val Acc:{best_acc:4f})# Load best model weightsmodel.load_state_dict(best_model_wts)returnmodel# Train single-channel CNNcriterionnn.CrossEntropyLoss()optimizer_singleoptim.Adam(single_channel_cnn.parameters(),lr0.001)print(Training Single-Channel CNN...)single_channel_cnntrain_model(single_channel_cnn,train_loader,val_loader,criterion,optimizer_single)# Train dual-channel CNN BiLSTM Attentionoptimizer_dualoptim.Adam(dual_channel_cnn_bilstm_att.parameters(),lr0.001)print(\nTraining Dual-Channel CNN BiLSTM Attention...)dual_channel_cnn_bilstm_atttrain_model(dual_channel_cnn_bilstm_att,train_loader,val_loader,criterion,optimizer_dual)2.4 结果可视化[titleResult Visualization]importcopy# Evaluate models on test setdefevaluate_model(model,test_loader):devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)model.to(device)model.eval()running_corrects0withtorch.no_grad():forsignals,images,labelsintest_loader:signalssignals.float().to(device)imagesimages.float().to(device)labelstorch.tensor(labels.argmax(axis1)).long().to(device)outputsmodel(signals,images)_,predstorch.max(outputs,1)running_correctstorch.sum(predslabels.data)accrunning_corrects.double()/len(test_loader.dataset)print(fTest Accuracy:{acc:.4f})# Evaluate single-channel CNNprint(Evaluating Single-Channel CNN...)evaluate_model(single_channel_cnn,test_loader)# Evaluate dual-channel CNN BiLSTM Attentionprint(\nEvaluating Dual-Channel CNN BiLSTM Attention...)evaluate_model(dual_channel_cnn_bilstm_att,test_loader)# Plot example STFT imagesplt.figure(figsize(12,4))foriinrange(4):plt.subplot(1,4,i1)plt.imshow(X_train_image[i],cmapviridis)plt.title(classes[y_train_class[i].argmax()])plt.axis(off)plt.show()运行步骤确保数据集路径正确将您的数据集放在datasets/case_western_reserve_university目录下。确保.npy文件存在并且格式正确。安装必要的库确保您已经安装了所需的库如numpy,pandas,torch,matplotlib,scipy等。您可以使用以下命令安装这些库pipinstallnumpy pandas torch matplotlib scipy scikit-learn运行代码直接运行上述完整的代码即可完成数据加载、预处理、模型构建、训练和评估。