保姆级教程：用PyTorch 2.0复现WDCNN轴承故障诊断模型（附CWRU数据集实战代码）

从零实现WDCNN轴承故障诊断PyTorch 2.0实战指南轴承作为机械设备的核心部件其健康状态直接影响整个系统的运行安全。传统故障诊断方法依赖专家经验而深度学习技术让自动化诊断成为可能。WDCNNWide Deep Convolutional Neural Network作为轴承故障诊断领域的经典模型凭借其独特的宽卷积核设计和抗噪能力在CWRU等标准数据集上表现出色。本文将手把手带你用PyTorch 2.0完整复现WDCNN模型从数据加载到模型部署每个环节都配有可运行的代码示例。1. 环境配置与数据准备工欲善其事必先利其器。在开始编码前我们需要搭建合适的开发环境。推荐使用Python 3.8和PyTorch 2.0以上版本这些版本在性能和功能上都有显著优化。安装核心依赖库pip install torch2.0.0 torchvision0.15.1 numpy1.23.5 pandas1.5.2 scikit-learn1.2.0 matplotlib3.7.1CWRU数据集是轴承故障诊断领域的基准数据集包含正常状态和多种故障类型内圈、外圈、滚动体故障在不同负载条件下的振动信号。数据集可从凯斯西储大学官网免费获取下载后建议按以下结构组织文件cwrudataset/ ├── 12k_drive_end/ │ ├── normal/ │ ├── inner_race/ │ ├── outer_race/ │ └── rolling_element/ └── 48k_fan_end/数据预处理是模型成功的关键。我们需要将原始振动信号转换为适合WDCNN输入的格式import numpy as np import pandas as pd from scipy.io import loadmat def load_cwru_data(file_path, segment_length1024): 加载CWRU数据文件并进行分段处理 :param file_path: .mat文件路径 :param segment_length: 每段信号长度 :return: 分段后的信号数组 (n_samples, segment_length) mat_data loadmat(file_path) key [k for k in mat_data.keys() if not k.startswith(__)][0] raw_signal mat_data[key].flatten() # 标准化处理 raw_signal (raw_signal - np.mean(raw_signal)) / np.std(raw_signal) # 分段处理 n_segments len(raw_signal) // segment_length segments np.zeros((n_segments, segment_length)) for i in range(n_segments): segments[i] raw_signal[i*segment_length : (i1)*segment_length] return segments提示CWRU数据采样频率为12kHz或48kHz建议统一降采样到12kHz以保证一致性。不同故障直径0.007英寸、0.014英寸等的数据应分开处理便于后续分析模型对不同严重程度故障的识别能力。2. WDCNN模型架构深度解析WDCNN的核心创新在于其第一层使用宽卷积核论文中为64这种设计可以直接处理原始振动信号无需复杂的特征工程。大卷积核能够捕获长周期的故障特征同时具有天然的降噪能力。让我们用PyTorch 2.0的nn.Module实现WDCNNimport torch import torch.nn as nn import torch.nn.functional as F class WDCNN(nn.Module): def __init__(self, in_channels1, num_classes10): super(WDCNN, self).__init__() # 第一层使用大卷积核是关键设计 self.conv1 nn.Sequential( nn.Conv1d(in_channels, 16, kernel_size64, stride16, padding24), nn.BatchNorm1d(16), nn.ReLU(), nn.MaxPool1d(kernel_size2, stride2) ) # 后续使用小卷积核堆叠 self.conv2 nn.Sequential( nn.Conv1d(16, 32, kernel_size3, padding1), nn.BatchNorm1d(32), nn.ReLU(), nn.MaxPool1d(kernel_size2, stride2) ) self.conv3 nn.Sequential( nn.Conv1d(32, 64, kernel_size3, padding1), nn.BatchNorm1d(64), nn.ReLU(), nn.MaxPool1d(kernel_size2, stride2) ) self.conv4 nn.Sequential( nn.Conv1d(64, 64, kernel_size3, padding1), nn.BatchNorm1d(64), nn.ReLU(), nn.MaxPool1d(kernel_size2, stride2) ) self.conv5 nn.Sequential( nn.Conv1d(64, 64, kernel_size3), nn.BatchNorm1d(64), nn.ReLU(), nn.MaxPool1d(kernel_size2, stride2) ) # 全连接层 self.fc nn.Sequential( nn.Linear(192, 100), # 注意输入维度需根据信号长度调整 nn.ReLU(), nn.Linear(100, num_classes) ) def forward(self, x): x self.conv1(x) x self.conv2(x) x self.conv3(x) x self.conv4(x) x self.conv5(x) x x.view(x.size(0), -1) # 展平 x self.fc(x) return x模型设计要点解析宽窄结合第一层宽卷积核64提取宏观特征后续小卷积核3捕捉细节特征降采样策略通过卷积stride和最大池化逐步降低时间维度特征图数量从16逐渐增加到64形成金字塔结构批归一化每层卷积后都加入BN层加速训练并提升泛化能力注意输入信号长度会影响最终全连接层的输入维度。对于1024点的信号经过五层降采样后特征维度为192。如果使用不同长度的信号需要相应调整全连接层的输入尺寸。3. 数据管道与训练流程高质量的数据管道能显著提升训练效率。我们使用PyTorch的Dataset和DataLoader构建高效的数据流from torch.utils.data import Dataset, DataLoader from sklearn.model_selection import train_test_split class BearingDataset(Dataset): def __init__(self, signals, labels, transformNone): self.signals signals self.labels labels self.transform transform def __len__(self): return len(self.labels) def __getitem__(self, idx): signal self.signals[idx] label self.labels[idx] if self.transform: signal self.transform(signal) # 增加通道维度 (1, signal_length) signal torch.FloatTensor(signal).unsqueeze(0) return signal, torch.LongTensor([label]) # 数据增强添加随机噪声 def add_noise(signal, noise_level0.05): noise np.random.normal(0, noise_level*np.std(signal), len(signal)) return signal noise # 划分训练集和测试集 def prepare_dataloaders(signals, labels, batch_size32, test_size0.2): X_train, X_test, y_train, y_test train_test_split( signals, labels, test_sizetest_size, stratifylabels ) train_dataset BearingDataset( X_train, y_train, transformadd_noise ) test_dataset BearingDataset(X_test, y_test) train_loader DataLoader( train_dataset, batch_sizebatch_size, shuffleTrue, num_workers4 ) test_loader DataLoader( test_dataset, batch_sizebatch_size, shuffleFalse, num_workers4 ) return train_loader, test_loader训练过程中有几个关键点需要注意学习率调度使用ReduceLROnPlateau在验证损失停滞时自动降低学习率早停机制防止过拟合混合精度训练利用PyTorch的AMP加速训练完整训练代码如下import torch.optim as optim from torch.cuda.amp import GradScaler, autocast from sklearn.metrics import accuracy_score def train_model(model, train_loader, val_loader, num_epochs50, lr0.001): device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lrlr) scheduler optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.5, patience5, verboseTrue ) best_val_loss float(inf) scaler GradScaler() # 混合精度训练 for epoch in range(num_epochs): model.train() train_loss 0.0 for signals, labels in train_loader: signals, labels signals.to(device), labels.to(device).squeeze() optimizer.zero_grad() # 混合精度训练 with autocast(): outputs model(signals) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() train_loss loss.item() * signals.size(0) # 验证阶段 model.eval() val_loss 0.0 all_preds, all_labels [], [] with torch.no_grad(): for signals, labels in val_loader: signals, labels signals.to(device), labels.to(device).squeeze() outputs model(signals) loss criterion(outputs, labels) val_loss loss.item() * signals.size(0) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) train_loss train_loss / len(train_loader.dataset) val_loss val_loss / len(val_loader.dataset) val_acc accuracy_score(all_labels, all_preds) print(fEpoch {epoch1}/{num_epochs} | fTrain Loss: {train_loss:.4f} | fVal Loss: {val_loss:.4f} | fVal Acc: {val_acc:.4f}) scheduler.step(val_loss) # 早停机制 if val_loss best_val_loss: best_val_loss val_loss torch.save(model.state_dict(), best_model.pth) patience 5 else: patience - 1 if patience 0: print(Early stopping triggered) break # 加载最佳模型 model.load_state_dict(torch.load(best_model.pth)) return model4. 模型评估与结果分析训练完成后我们需要全面评估模型性能。除了准确率故障诊断任务还需要关注各类别的精确率、召回率和F1分数混淆矩阵分析不同信噪比下的鲁棒性from sklearn.metrics import classification_report, confusion_matrix import seaborn as sns import matplotlib.pyplot as plt def evaluate_model(model, test_loader): device torch.device(cuda if torch.cuda.is_available() else cpu) model.eval() all_preds, all_labels [], [] with torch.no_grad(): for signals, labels in test_loader: signals, labels signals.to(device), labels.to(device).squeeze() outputs model(signals) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) # 分类报告 print(Classification Report:) print(classification_report(all_labels, all_preds)) # 混淆矩阵 cm confusion_matrix(all_labels, all_preds) plt.figure(figsize(10, 8)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues) plt.xlabel(Predicted) plt.ylabel(Actual) plt.title(Confusion Matrix) plt.show() return all_preds, all_labels典型评估结果示例指标正常状态内圈故障外圈故障滚动体故障精确率99.2%97.8%98.5%96.3%召回率98.7%96.5%99.1%95.8%F1分数98.9%97.1%98.8%96.0%WDCNN在CWRU数据集上的表现通常能达到97%以上的总体准确率特别是对于早期故障小直径缺陷的检测能力明显优于传统方法。模型的第一层大卷积核使其对噪声具有较强的鲁棒性在信噪比降至10dB时仍能保持90%以上的识别准确率。5. 实战技巧与常见问题在实际项目中应用WDCNN时有几个实用技巧值得分享输入信号长度选择1024点是最常用的长度平衡了计算效率和特征完整性对于低速设备可能需要更长的信号以包含完整的旋转周期可以使用nn.AdaptiveMaxPool1d替代固定尺寸的池化适应不同长度输入处理类别不平衡# 计算类别权重 from sklearn.utils.class_weight import compute_class_weight class_weights compute_class_weight( balanced, classesnp.unique(train_labels), ytrain_labels ) class_weights torch.FloatTensor(class_weights).to(device) criterion nn.CrossEntropyLoss(weightclass_weights)模型轻量化减少中间层的通道数如从64减至32使用深度可分离卷积替代标准卷积量化训练后的模型跨设备泛化# 域适应技巧在最后一层前加入梯度反转层 class GradientReversalFunction(torch.autograd.Function): staticmethod def forward(ctx, x, alpha): ctx.alpha alpha return x.view_as(x) staticmethod def backward(ctx, grad_output): return grad_output.neg() * ctx.alpha, None class GradientReversal(nn.Module): def __init__(self, alpha1.0): super().__init__() self.alpha alpha def forward(self, x): return GradientReversalFunction.apply(x, self.alpha)常见问题排查问题验证准确率波动大检查学习率是否过高batch size是否过小解决减小学习率增大batch size或使用学习率预热问题训练损失下降但验证损失上升检查是否过拟合解决增加数据增强添加Dropout层或使用更早的停止点问题模型对所有样本预测同一类别检查类别是否极度不平衡最后一层偏置是否过大解决使用类别加权损失检查初始化方式6. 进阶优化方向基础WDCNN模型已经表现不错但仍有优化空间多传感器融合class MultiSensorWDCNN(nn.Module): def __init__(self, in_channels3, num_classes10): # 3通道输入 super().__init__() self.conv1 nn.Sequential( nn.Conv1d(in_channels, 16, kernel_size64, stride16, padding24), nn.BatchNorm1d(16), nn.ReLU(), nn.MaxPool1d(2, 2) ) # 其余层与原始WDCNN相同时频联合特征class TimeFrequencyWDCNN(nn.Module): def __init__(self, num_classes10): super().__init__() # 时域路径 self.time_path WDCNN(in_channels1, num_classesnum_classes) # 频域路径 self.freq_conv nn.Sequential( nn.Conv1d(1, 16, kernel_size64, stride16, padding24), nn.BatchNorm1d(16), nn.ReLU(), nn.MaxPool1d(2, 2) ) # 合并两路特征 self.fc nn.Linear(384, num_classes) # 假设每路输出192维注意力机制增强class AttentionBlock(nn.Module): def __init__(self, channels): super().__init__() self.query nn.Conv1d(channels, channels//8, 1) self.key nn.Conv1d(channels, channels//8, 1) self.value nn.Conv1d(channels, channels, 1) self.gamma nn.Parameter(torch.zeros(1)) def forward(self, x): B, C, L x.shape q self.query(x).view(B, -1, L) # (B, C//8, L) k self.key(x).view(B, -1, L) # (B, C//8, L) v self.value(x).view(B, C, L) # (B, C, L) attn torch.bmm(q.transpose(1, 2), k) # (B, L, L) attn F.softmax(attn, dim-1) out torch.bmm(v, attn) return self.gamma * out x在线学习与增量训练def incremental_train(model, new_loader, old_loaderNone, epochs10): # 保存旧数据的重要样本 if old_loader: exemplars select_exemplars(old_loader, model) # 结合新数据和旧样本 combined_loader combine_loaders(new_loader, exemplars) # 微调模型 optimizer optim.SGD(model.parameters(), lr0.001, momentum0.9) for epoch in range(epochs): model.train() for inputs, labels in combined_loader: # 标准训练流程 ...轴承故障诊断是一个快速发展的领域WDCNN作为经典模型为后续研究奠定了重要基础。在实际工业应用中还需要考虑计算资源限制、实时性要求等因素可能需要简化模型结构或进行量化压缩。