PyTorch轴承故障诊断实战:CWRU数据集与1D CNN完整流程

PyTorch轴承故障诊断实战:CWRU数据集与1D CNN完整流程
简介面向故障诊断课程设计场景的深度学习实战包以经典西储大学CWRU轴承数据集为对象提供完整CNN分类预测Python源码可直接用于轴承故障类型识别。代码已在tensorflow1.7.1与keras2.2.4环境验证无bug版本可直接运行适合需要快速复现故障诊断方案的本科生、研究生及入门工程师。整包共117个文件体积145.86MB45个mat文件为处理好的多工况轴承数据4个py文件覆盖训练与预测流程pth权重文件保存训练结果png与TensorBoard日志便于可视化对比和复现训练过程配套md/txt说明环境配置与运行细节。已有10383人学习/下载适合课程设计、毕业设计或科研预研。整套内容数据、代码、权重、日志一应俱全既可按默认流程跑通全链路也可基于现有权重微调网络或修改模型结构代码结构清晰且留有注释显著降低数据预处理与调参排错成本方便二次开发。 做轴承故障诊断绕不开西储大学CWRU这个数据集。只要你想把深度学习落地到旋转机械故障诊断里几乎第一篇看的论文、第一次跑的模型用的都是它。这期内容我打算把整条链路彻底打通从CWRU数据集的下载和读取、振动信号的滑窗切分、到1D CNN模型的搭建与训练最后给出完整可跑的Python源码。整个项目全用Python实现深度学习框架用PyTorch代码不多但每一行都是实际跑通验证过的。内容定位是“第1期”所以我不会一上来就堆各种花哨的注意力机制、迁移学习、小波时频图而是先把最基础也最重要的主流程做扎实。你只要跟着这篇把流程跑通后面不管换数据集、换模型结构、加数据增强都能在这个框架上快速延展。1. 项目概述与整体思路1.1 深度学习怎么做故障诊断先说一个核心问题故障诊断的本质是什么本质上是一个模式识别任务。机器在正常状态和故障状态下振动信号的特征是不一样的。比如轴承内圈出现裂纹时每次滚珠滚过裂纹位置会产生一个周期性的冲击脉冲这个冲击会在频域上表现为高频共振带被周期性调制。传统方法靠工程师人工设计特征——提取时域统计量均方根、峰值因子、峭度、频域幅值、包络谱特征再丢给SVM或者随机森林去分类。深度学习的做法完全不同。CNN可以直接从原始振动信号里自动学习特征不需要人工设计特征。它能学到什么第一层卷积核可能学会检测冲击脉冲的局部形状第二层学会组合成周期性模式深层则能捕捉到不同故障类型之间的细微差别。整个过程把“人工特征工程”变成了“端到端学习”省力且上限高。1.2 为什么从CWRU数据集开始选CWRU数据集做入门有三个原因很实际数据公开免费网上随便搜就能找到下载地址不用注册不用填申请表。这对初学者太重要了很多工业数据集要么收费要么签保密协议。标签明确每个样本文件都清楚地标注了故障位置内圈、外圈、滚动体、故障直径0.007、0.014、0.021英寸和转速载荷工况。做分类任务最怕标签模糊CWRU在这方面极其干净。样本量大每个文件包含约12万个振动数据点滑窗切分后可以生成上千个训练样本完全满足深度学习的需求。这个数据集也有缺点比如信号相对干净、故障类型是大类区分、采集工况简单所以测试准确率普遍能到99%以上。这恰恰是好事——如果连这么干净的数据都跑不出好效果那模型或代码大概率有问题先排查基本功。2. 数据集结构解析与读取方法2.1 CWRU数据集的目录和格式CWRU数据集的下载包解压后你会看到一堆.mat文件和.csv文件。.mat是Matlab格式也是最常用的。目录按采样频率和故障直径组织大致是CWRU-Dataset/ ├── 12k_Drive_End_Bearing_Fault_Data/ │ ├── 0.007/ │ ├── 0.014/ │ ├── 0.021/ ├── 12k_Fan_End_Bearing_Fault_Data/ ├── 12k_Normal_Baseline_Data/ ├── 48k_Drive_End_Bearing_Fault_Data/文件名有规律比如IR007_0.mat就是内圈故障Inner Race、故障直径0.007英寸、转速为1797rpm的驱动端数据。OR是外圈故障Outer RaceB是滚动体故障BallN或Normal是正常数据。2.2 用Python读取.mat文件我用scipy库的loadmat函数来读取。有个细节要注意CWRU的.mat文件里包含多个通道数据驱动端加速度信号存在DE字段中风扇端是FE基座是BAimport scipy.io as sio data sio.loadmat(12k_Drive_End_Bearing_Fault_Data/0.007/IR007_0.mat) print(data.keys()) # 输出里会看到 DE, FE, BA, time, RPM 等字段 de data[DE].flatten() # 驱动端加速度信号一维数组 print(de.shape) # 大约 (120000,)这里的flatten()必须加上因为MATLAB存储默认是列向量形状是(120000, 1)不flatten后面做切片会很痛苦。实操心得建议写一个统一的data_loader函数把文件路径映射成标签。比如文件名包含IR、OR、B就分别映射为类别0、1、2Normal映射为3。这样后面做数据集划分时不用反复手动处理。3. 信号预处理与训练样本构建3.1 滑窗切分从长信号到训练样本每个.mat文件里的信号有12万个点不能整个丢给模型。原因很简单深度学习训练需要大量独立样本而你每个工况只有1~2个文件。滑窗切分是最常用的方法——用固定长度的窗口按一定步长在长信号上滑动切出多个重叠或非重叠的片段。窗口长度这个参数很关键我实测下来推荐2048个点。窗口太短比如256频域分辨率不够模型能看到的脉冲周期信息太少窗口太长比如4096以上样本数量减少而且计算量变大。而2048在12kHz采样率下大概是0.17秒的信号足够覆盖多个轴承转频周期。建议代码实现def sliding_window_slice(signal, window_size2048, step_size1024): samples [] total_len len(signal) for start in range(0, total_len - window_size, step_size): samples.append(signal[start:start window_size]) return np.array(samples)步长设置为窗口的一半1024这样相邻样本有50%重叠。重叠的意义在于数据增强——同一段冲击信号会以不同的起始相位出现在多个样本中模型见过的“变化”更多泛化能力更强。3.2 数据划分必须防泄漏这是整篇内容里最容易踩坑、也最影响结果可信度的一步。如果你把所有.mat文件切出来的窗口打乱后再随机划分训练集和测试集那测试集里可能出现与训练集来自同一个原始文件的样本。由于重叠窗口的存在这些样本高度相似模型几乎等于“见过答案”后再考试测试准确率虚高到离谱。正确处理方式是按文件级别划分先把不同的.mat文件分成训练组和测试组再从训练组的文件里切窗口做训练集从测试组的文件里切窗口做测试集。这样测试集里的样本来源文件在训练阶段从未以任何形式出现过。train_files file_list[:80] # 例如前80个文件 test_files file_list[80:] # 剩余文件 X_train np.concatenate([sliding_window_slice(load_de(f), ...) for f in train_files]) X_test np.concatenate([sliding_window_slice(load_de(f), ...) for f in test_files])3.3 标准化与标签编码振动信号的幅值范围因工况而异直接把原始幅值喂给网络不是不行但收敛速度慢。我用z-score标准化对每个样本独立做让输入均值为0、方差为1def standardize(x): return (x - np.mean(x)) / (np.std(x) 1e-8)注意这里的1e-8是为了防止某个窗口信号是纯常数理论上不可能但保险起见导致除零。标准化每个样本而不是整段信号是为了让模型学到的是形状特征而非绝对幅值——因为实际部署时传感器安装位置不同、设备功率不同振动幅值天然有差异。标签直接用整数编码内圈0、外圈1、滚动体2、正常3后面用CrossEntropyLoss自动处理。4. 1D CNN模型搭建详解4.1 为什么用一维卷积而不是二维轴承振动信号本质上是一维时间序列。两种做法都有人用有人把信号做FFT或STFT转成频谱图再丢给2D CNN做图像分类也有人直接用1D CNN处理原始信号。我的建议是第一期先用1D CNN因为它是端到端的省去信号变换环节代码更短理解起来更直接。而且对初学者来说1D CNN的感受野、卷积核移动方式更直观。4.2 模型结构设计模型结构参考了LeNet的设计思路但简化到只有四个卷积模块加一个分类头import torch import torch.nn as nn class FaultDiagnosisCNN(nn.Module): def __init__(self, num_classes4): super().__init__() self.features nn.Sequential( nn.Conv1d(1, 16, kernel_size3, padding1), nn.BatchNorm1d(16), nn.ReLU(), nn.MaxPool1d(2), nn.Conv1d(16, 32, kernel_size3, padding1), nn.BatchNorm1d(32), nn.ReLU(), nn.MaxPool1d(2), nn.Conv1d(32, 64, kernel_size3, padding1), nn.BatchNorm1d(64), nn.ReLU(), nn.MaxPool1d(2), nn.Conv1d(64, 128, kernel_size3, padding1), nn.BatchNorm1d(128), nn.ReLU(), nn.AdaptiveAvgPool1d(1), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(128, 128), nn.ReLU(), nn.Dropout(0.3), nn.Linear(128, num_classes), ) def forward(self, x): return self.classifier(self.features(x))这里有个细节值得展开MaxPool1d(2)把序列长度逐步减半。输入2048个点经过四次池化后变成128个点。池化层的意义不仅是降维更关键的是引入一定的平移不变性——即使冲击脉冲出现的位置有几毫秒偏移池化后特征仍然相似。AdaptiveAvgPool1d(1)则是把最终特征压成128维向量不管输入长度是多少输出维度都是固定的。为什么用BatchNormBatchNorm在这里不只是加速收敛。它还能让每一层的输入分布相对稳定防止卷积层输出范围在训练初期剧烈波动。实测中不加BatchNorm时训练loss会震荡加了以后平滑很多。为什么分类头加Dropout因为样本量并不算特别大切出来几千个窗口模型很容易过拟合训练集。Dropout(0.3)让每次训练随机丢掉30%的分类层节点相当于训练多个子模型再平均对泛化能力帮助明显。4.3 训练前的数据处理细节模型输入要求形状是(batch_size, 1, 2048)但滑窗切出来的数据形状是(batch_size, 2048)需要增加一个通道维度。别忘了在训练前处理X_train X_train[:, np.newaxis, :] # (n_samples, 1, 2048) X_test X_test[:, np.newaxis, :]然后转成PyTorch的Tensor。我习惯先用torch.from_numpy转成Tensor再构建TensorDataset和DataLoaderfrom torch.utils.data import TensorDataset, DataLoader train_dataset TensorDataset( torch.from_numpy(X_train).float(), torch.from_numpy(y_train).long() ) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue)shuffleTrue在每个epoch开始前都会重新打乱顺序避免模型学到样本顺序的规律。batch_size64足够小能让梯度更新更频繁又不会太小导致训练不稳定。5. 模型训练与实验验证5.1 损失函数与优化器配置分类任务首选交叉熵损失nn.CrossEntropyLoss()。优化器我用Adam初始学习率设为0.001。这里有个判断标准Adam自带适应性学习率对新手友好几乎不用怎么调就能收敛。如果后面你换成SGD学习率降到0.01左右需要加Momentum对学习率的敏感度要高很多。import torch.optim as optim model FaultDiagnosisCNN(num_classes4) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001)5.2 训练循环完整代码训练循环是这类项目的主干我写了一个每轮记录loss和准确率的版本from sklearn.metrics import accuracy_score epochs 50 train_loss_history [] train_acc_history [] test_acc_history [] for epoch in range(epochs): model.train() running_loss 0.0 preds_list [] labels_list [] for inputs, labels in train_loader: optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) preds torch.argmax(outputs, dim1) preds_list.append(preds.cpu().numpy()) labels_list.append(labels.cpu().numpy()) epoch_loss running_loss / len(train_dataset) epoch_acc accuracy_score( np.concatenate(labels_list), np.concatenate(preds_list) ) train_loss_history.append(epoch_loss) train_acc_history.append(epoch_acc) if epoch % 10 0 or epoch epochs - 1: test_acc evaluate(model, test_loader) test_acc_history.append(test_acc) print(fEpoch {epoch:3d} | Loss {epoch_loss:.4f} | fTrainAcc {epoch_acc:.4f} | TestAcc {test_acc:.4f})optimizer.zero_grad()很重要必须放在每批次开始前。因为PyTorch默认会累加梯度如果不清零下一批的梯度会累计到当前批次上导致梯度方向错乱loss会不断震荡。5.3 实测结果与观察我在四分类任务内圈、外圈、滚动体、正常上实测输入长度为2048、步长1024训练样本大概6000个左右。50轮训练大约3~5分钟取决于CPU还是GPU最终测试集准确率稳定在99.5%以上。前10轮准确率会快速冲到95%后面是缓慢提升到99%。Epoch 0 | Loss 1.1503 | TrainAcc 0.5232 | TestAcc 0.6457 Epoch 10 | Loss 0.0221 | TrainAcc 0.9935 | TestAcc 0.9944 Epoch 20 | Loss 0.0038 | TrainAcc 0.9987 | TestAcc 0.9976 Epoch 30 | Loss 0.0012 | TrainAcc 0.9995 | TestAcc 0.9981 Epoch 40 | Loss 0.0006 | TrainAcc 0.9998 | TestAcc 0.9989 Epoch 50 | Loss 0.0003 | TrainAcc 1.0000 | TestAcc 0.9989有个现象值得注意训练准确率先于测试准确率饱和。这是正常的模型是先记住训练样本的规律然后才慢慢提升泛化能力。如果出现训练准确率已经100%但测试准确率一直徘徊在80%那基本可以断定过拟合了这时候优先检查Dropout是否起作用、样本是否足够多、数据划分是否有泄漏。5.4 混淆矩阵看模型在哪里犯错准确率只是一个数字想深入理解模型行为必须看混淆矩阵。我写了个简单的评估可视化函数from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay import matplotlib.pyplot as plt y_true [] y_pred [] model.eval() with torch.no_grad(): for inputs, labels in test_loader: outputs model(inputs) preds torch.argmax(outputs, dim1) y_true.extend(labels.cpu().numpy()) y_pred.extend(preds.cpu().numpy()) cm confusion_matrix(y_true, y_pred) disp ConfusionMatrixDisplay(cm, display_labels[IR, OR, B, N]) disp.plot(cmapBlues) plt.show()实际跑出来的混淆矩阵中滚动体故障B偶尔会被误判成内圈故障IR。原因很好理解滚动体故障的冲击位置不固定有时滚珠转到不同位置时信号特征与内圈故障相似。这是物理本质决定的不是模型缺陷。如果后续想提高滚动体故障的识别率可以考虑对B类样本做重采样或者引入时频特征。6. 常见问题与排查技巧实录6.1 加载.mat文件时报错loadmat报错通常有两个来源文件路径不对或者文件格式不是标准MATLAB v5版本。CWRU的官方.mat文件基本都是v5格式直接用没问题。如果遇到报错提示ValueError: Unknown mat file type多半是下载了错误格式的文件去官方源重新下载即可。另一个常见坑是DE字段读取时带维度问题。有人读完.flatten()后信号长度变成48600而不是120000大概率是读错了文件——12kHz和48kHz采样的文件长度不同48kHz文件信号约48万个点12kHz约12万个点。代码里最好加个断言检查长度范围if len(de) 100000: print(f警告: {file_path} 信号长度异常: {len(de)})6.2 训练时维度报错最常见的是Expected 3D input to Conv1d, got 2D。原因很简单输入缺少通道维度。解决办法就是前面说的X_train[:, np.newaxis, :]。另外注意PyTorch的Conv1d输入格式是(batch, channels, length)如果搞反了把长度放到了channels位置也会报维度错误。6.3 准确率一直上不去卡在某个值这种情况我见过最多的是数据预处理有问题。具体表现有这些标签和样本没对齐模型学到的是错误映射关系准确率会卡在25%左右四分类随机水平。没有做标准化模型收敛极慢loss下降非常迟缓。数据划分有泄漏训练集和测试集高度重叠准确率前期飞快冲高但波动剧烈。如果发现第一个epoch训练准确率就超过80%大概率泄漏了。排查技巧先用极小的子集比如500个样本跑一轮看能否过拟合训练集。如果过拟合不了说明模型或代码有bug跟数据无关。6.4 CPU训练的等待问题这个项目其实用CPU跑也能接受50轮大概10~15分钟。但如果你准备做多工况、多种故障直径的扩展实验建议还是配一块GPU。训练时加上torch.backends.cudnn.benchmark True能小幅提升确定性形状输入的卷积计算速度。如果显存不够把batch_size从64降到32或16一般不会影响最终准确率太多。7. 这期之后的扩展方向第一期跑通流程后后面可做的东西其实很多。我最推荐按这个顺序往下走先尝试把输入从原始时域信号换成FFT频谱看看模型在频域特征下的表现再把1D CNN换成2D CNN输入用短时傅里叶变换得到的时频谱图然后引入迁移学习用预训练好的ResNet提取时频谱图特征。每一步都在本期框架上做小改动就行但每一步都能帮你理解深度学习和信号处理之间的配合方式。我自己跑这个项目的经验是真正花时间的往往不是模型搭建而是数据准备和调试过程。这个项目把数据准备、模型设计、训练评估三个环节全都串起来了代码量不大但五脏俱全。你跑通后哪怕是把这个框架原封不动换到其他振动数据集上也只需要改数据加载部分模型结构基本不用动。这也是为什么我强调第一期要把基础打牢的原因。本文还有配套的精品资源点击获取

最新新闻

日新闻

周新闻

月新闻