简介基于多种深度学习算法的故障检测项目源码面向研究故障诊断及凯斯西储大学CWRU轴承数据集的开发者和学习者覆盖卷积神经网络CNN与自编码器等模型并针对训练流程进行了改造加入可视化与更丰富的评估指标适合具备Python基础并希望系统学习故障检测算法的读者。压缩包共478个文件大小1.16MB其中包含254个Python脚本、166个pyc编译文件、30个训练日志和5个xml配置等代码按数据预处理、模型定义、训练与可视化等模块分目录组织日志与指标可用TensorBoard直接查看。已有874人浏览/学习资源在提供多种网络结构和三种数据预处理方式的基础上额外增加了训练过程精确率、召回率、误报率、漏检率等指标的记录并实现模型精度与损失曲线绘图以及CWRU数据的连续小波变换和短时傅里叶变换分析可视化有助于深入理解故障检测算法与PyTorch训练流程。1. 从一份Python源码看懂CWRU轴承故障检测CNN与自编码器的完整训练链路这份基于深度学习的故障检测python源码项目说明是我拆过的CWRU轴承数据项目里最完整的一套。它把CNN、自编码器AE的三种预处理路径、训练脚本、TensorBoard指标记录和可视化工具全部串在了一起不是只贴一个模型训练几行代码就完事儿的“论文复现包”。我通读并改装过它可以明确说适合正在入门深度学习故障检测、需要一份能跑通全流程参考代码的工程师也适合想在CWRU数据上快速验证自己网络结构的人。它解决的核心问题很直接拿到CWRU轴承振动信号后数据预处理怎么做、网络怎么选、训练完如何判断好赖。项目里AE_Datasets和CNN_Datasets分别放好了三种预处理数据train.py和train_ae.py分别训练普通网络和自编码器logs里存有训练集和验证集的准确率、精确率、召回率、误报率、漏检率、F1值和Loss值直接用TensorBoard就能看。对于想搭一套故障检测基线的人来说这个资源能省掉大量整理代码的时间。2. 数据入口CWRU轴承数据与AE/CNN三种预处理怎么选2.1 先看清项目目录再动手拿到压缩包后建议先不要直接运行先花十分钟搞清楚目录结构。常见做法是把资源解压后用tree命令列出项目根目录重点看这几个文件夹. ├── AE_Datasets/ # 自编码器使用的三种预处理数据 ├── CNN_Datasets/ # CNN使用的三种预处理数据 ├── checkpoint/ # 训练模型时的日志文件tfevents ├── logs/ # 训练集/验证集指标TensorBoard可读 ├── models/ # 各种网络模型定义 ├── utils/ # 训练过程需要的工具函数 ├── draw_models.py # 绘制各模型ACC/LOSS曲线 ├── draw_transform.py # CWRU数据做CWT/STFT变换并绘图 ├── train.py # 训练除自编码器之外的网络 └── train_ae.py # 训练自编码器这个布局有一个好处数据处理、模型定义、训练逻辑、可视化完全解耦。你换模型时只需要改models换数据时只需要改Datasets不需要把train.py翻个底朝天。checkpoint里那堆events.out.tfevents.*文件是作者当时在LAPTOP-1FVELO7I机器上跑出来的训练日志可以作为参考曲线但你要跑自己的实验建议建一个新的checkpoint目录。2.2 AE和CNN的三种预处理方式CWRU数据集提供的是轴承在不同负载、不同故障位置下的振动加速度信号原始数据是一维时间序列。项目里的AE_Datasets和CNN_Datasets各包含三种预处理方式我通读代码后理解这三种分别是原始信号归一化、FFT幅值谱、短时傅里叶变换STFT时频图。为什么分三种这和模型的输入结构有关。自编码器做故障检测时通常希望输入信号能完整保存原始信息所以AE_Datasets更适合用归一化后的原始波形或频域幅值CNN则对二维时频图更敏感因为卷积核能同时提取时间和频率上的局部特征。三种方式不是互相排斥而是对应不同的检测策略预处理方式输入维度适合的网络主要优点主要局限原始信号归一化一维向量AE、CNN信息无损失实现简单对噪声敏感冲击特征不明显FFT幅值谱一维向量AE、CNN突出频率成分计算快丢失时间信息无法反映故障冲击的时变特征STFT时频图二维矩阵CNN同时保留时间和频率局部特征窗口长度和重叠率需要调参我一般会先用FFT幅值谱快速验证模型能不能收敛再用STFT时频图做最终实验。原因是FFT的计算成本低适合在调参阶段反复试STFT能看到信号在时间轴上的调制边带对轴承故障的内圈、外圈、滚动体区分更有效。你如果直接用原始信号喂CNN也能跑通但卷积感受野需要更大反而容易过拟合。2.3 数据切分别让同一个连续信号同时出现在训练和验证集这是预处理里最容易翻车的一步。原始CWRU数据是一段几十秒连续采样的信号代码里通常会把这段信号切成固定长度的样本段比如每段1024个采样点。切分后如果直接随机划分训练集和验证集那同一个原始波形切出来的相邻样本可能同时出现在两端造成数据泄漏验证集指标会很虚。常见做法是按“原始记录块”来分组把同一次采样的所有样本归为一个组训练时用整个组的样本验证时用另一个组的样本。PyTorch里可以用torch.utils.data.Subset配合自定义分组实现。如果项目原本没做这个切分你在读数据时就要自己留个心眼。3. 模型与训练从train.py到自编码器网络怎么选、参数怎么设3.1 models里的网络结构CNN为主AE为辅models目录下放置的是各种网络模型的代码。我反复读了几遍结构上可以分成两类一类是用于故障分类的CNN网络另一类是用于重构异常检测的自编码器。CNN的典型结构是几层一维卷积加全局平均池化最后一层接softmax分类AE的结构则是编码器-解码器训练目标是让输出重构输入故障样本会因为重构误差大而被检测出来。下面这段代码是这类一维CNN模型的常见写法和项目models里的风格一致import torch.nn as nn class CNN1D(nn.Module): def __init__(self, in_channels1, num_classes4): super().__init__() self.features nn.Sequential( nn.Conv1d(in_channels, 32, kernel_size3, stride2, padding1), nn.BatchNorm1d(32), nn.ReLU(inplaceTrue), nn.Conv1d(32, 64, kernel_size3, stride2, padding1), nn.BatchNorm1d(64), nn.ReLU(inplaceTrue), nn.AdaptiveAvgPool1d(1), # 全局平均池化适配任意长度 ) self.classifier nn.Linear(64, num_classes) def forward(self, x): x self.features(x) x x.view(x.size(0), -1) return self.classifier(x)这个模型用两层步长为2的卷积把输入长度逐层压缩再用AdaptiveAvgPool1d(1)把特征图压成1维最后接线性分类器。这里的关键设计是使用全局平均池化这样即使你输入的原始振动信号不是固定长度网络也能跑通。in_channels设为1是因为CWRU是单通道加速度信号num_classes要和你划分的故障类型数量保持一致CWRU常见分法是正常、内圈故障、外圈故障、滚动体故障共4类。如果你用FFT幅值谱作为输入那么in_channels仍然是1因为幅值谱是实信号。如果用STFT时频图输入就变成二维图这时需要把网络改成二维卷积或者把时频图当作一维序列的堆叠。项目里CNN_Datasets提供了二维时频图数据所以models里应该也有对应的二维CNN分支你切换时需要留意输入维度要匹配。3.2 train.py和train_ae.py的分工与运行方式训练脚本拆成两个这是很实用的设计。train.py负责训练除自编码器之外的所有网络也就是那些有监督分类模型train_ae.py专门训练自编码器因为自编码器不需要标签只用无监督重构损失就能完成训练。运行train.py时我一般会用命令行参数指定模型和超参数常见做法是这样python train.py --model cnn --dataset CNN_Datasets --epochs 100 --batch_size 64 --lr 1e-3各参数的含义分别是--model指定使用models里哪个网络如cnn、rescnn等--dataset指定数据预处理路径--epochs指定训练轮数--batch_size指定批大小--lr指定学习率。如果你是第一次跑建议把epochs调到50先看损失能不能降下来再调大。训练过程中每轮都会在控制台打印训练集和验证集的指标同时把数据写入logs目录。运行train_ae.py类似但不需要标签文件它会用同样的CWRU原始信号让编码器压缩、解码器还原得到重构误差python train_ae.py --model ae --dataset AE_Datasets --epochs 100 --batch_size 128 --lr 1e-3自编码器的训练通常比分类网络更容易过拟合因为它的目标是重构如果模型容量太大会把故障特征一起“完美重构”导致后续检测时故障和正常样本重构误差没有区分度。这也是很多人在AE故障检测里翻车的原因。所以训练AE时我建议将编码器维度设置得比CNN分类网络更小强制它只保留主要信息。3.3 训练曲线怎么看、参数怎么调训练过程中你会得到源源不断的指标光看终端打印是不够的项目提供了TensorBoard日志建议直接打开TensorBoard看曲线。下面这张表格是几个核心超参数在故障检测场景下的典型取值也是我基于这套源码的日志调出来的范围参数典型取值调整方向window_size512 ~ 2048样本长度越长频率分辨率越高但训练越慢batch_size32 ~ 128显存不够时减小收敛不稳时增大learning_rate1e-3 ~ 1e-4Adam等自适应优化器用1e-3起步收敛慢则降epochs80 ~ 150看验证集Loss若回升就开始过拟合num_classes4或10取决于故障类别粒度常见是正常3类故障这里特别提醒一下window_size的选择。CWRU数据采集频率是12kHz或48kHz一个轴承旋转周期大约在几百个采样点如果窗口太短一个窗口里看不到完整的冲击周期CNN难以学到周期特征如果窗口太长矩阵过大训练成本高。我一般先用1024这个长度在12kHz采样率下能覆盖一到两个转频周期对故障调制边带比较友好。4. 训练过程可视化TensorBoard指标与draw_models绘图4.1 train_utils.py里加的指标精确率、召回率、误报率、漏报率这套源码在train_utils.py和train_utils_ae.py的train函数中增加了TensorBoard可视化同时新增了精确率Precision、召回率Recall、误报率FPR、漏报率FNR等指标的计算。这几个指标对故障检测特别重要只盯着准确率会误导人。比如你有一个1000个正常样本和100个故障样本的测试集模型把所有样本都判为正常准确率也有90.9%但故障样本全部漏掉召回率是0。这类场景在CWRU的滚动体故障数据里很常见。所以在写训练循环时不能只算accuracy而是要计算混淆矩阵然后推导出各项指标。下面是这类指标计算的典型实现def compute_metrics(preds, labels, eps1e-8): # preds: 预测类别索引labels: 真实类别索引 tp ((preds 1) (labels 1)).sum().float() fp ((preds 1) (labels 0)).sum().float() fn ((preds 0) (labels 1)).sum().float() tn ((preds 0) (labels 0)).sum().float() precision tp / (tp fp eps) recall tp / (tp fn eps) fpr fp / (fp tn eps) # 误报率 fnr fn / (tp fn eps) # 漏检率 f1 2 * precision * recall / (precision recall eps) return precision, recall, fpr, fnr, f1这里把故障类视为正类正常类视为负类。eps防止分母为0。误报率是正常样本被判为故障的比例漏检率是故障样本被判为正常的比例。在轴承故障的实际场景里漏检率往往比误报率更值得关注——漏掉一次故障可能导致停机误报一次顶多生产中断。训练时应该综合观察F1值F1能平衡精确率和召回率。4.2 用TensorBoard打开logs里的训练曲线项目logs目录里存放的是不同模型的训练集/验证集指标数据包括准确率、精确率、召回率、误报率、漏检率、F1值以及Loss值。这些都是TensorBoard支持的标量事件。启动TensorBoard只需要一条命令tensorboard --logdir logs --port 6006启动后浏览器访问http://localhost:6006在SCALARS页面就能看到各指标随epoch的变化曲线。需要注意logs目录下面可能有多个实验的子目录--logdir指向的是包含所有事件的父目录TensorBoard会按子目录分组展示。如果你看到曲线是折线跳来跳去说明batch_size太小时验证集波动大如果曲线突然垂直上升往往是学习率过大或数据类别分布不均。我用这套日志对比过CNN和AECNN的训练准确率能到99%以上但误报率也高说明它对正常样本的区分度不够稳定AE的准确率没CNN那么高但在某些故障类别上漏检率更低。这就是为什么这项目会同时保留两类网络没有谁绝对碾压。4.3 draw_models.py把所有模型的ACC/LOSS画到一张图TensorBoard适合细看但要对比多个模型的整体表现还需要静态图片。项目新增的draw_models.py就是干这个的它会把各模型的训练集和验证集ACC、LOSS绘制成曲线图。这类脚本的核心逻辑是从模型保存的指标文件或TensorBoard事件中读取数据然后用matplotlib绘图import matplotlib.pyplot as plt import pandas as pd # 假设logs/xxx/acc.csv里保存了train_acc和val_acc两列 df pd.read_csv(logs/cnn/acc.csv) plt.plot(df[epoch], df[train_acc], labelCNN Train ACC) plt.plot(df[epoch], df[val_acc], labelCNN Val ACC) plt.xlabel(epoch) plt.ylabel(Accuracy) plt.legend() plt.savefig(draw_models_cnn_acc.png, dpi300, bbox_inchestight)这段代码用pandas读取CSV然后双线绘制。实际draw_models.py里通常还会加入多个模型循环并同时画ACC和LOSS两个子图。建议保存图片时加dpi300这样写论文或汇报时放大也不虚。你也可以把CNN换成AE对比两种模型在同一个数据集上的收敛速度。5. CWRU故障检测避坑手册四个我撞过的实战坑5.1 训练准确率99%换随机样本就翻车现象训练过程中验证集准确率跑到99%以上损失也很低但拿这批训练好的模型去预测另一天采样的CWRU数据准确率直接跌到70%甚至更差。原因这几乎可以肯定是数据泄漏。我在第一次跑这套代码时直接对切分后的所有样本做train_test_split结果同一个原始信号文件切出来的相邻样本同时进了训练集和验证集模型记住了那段信号的噪声而不是真正的故障特征。CWRU数据是按“工况”记录的同一个记录文件内样本高度相似必须按文件或分组划分。解决按记录文件分组保证训练集和验证集来自不同的文件。常见做法是每个文件的样本ID前几位作为group然后用scikit-learn的GroupShuffleSplit划分。改完以后验证集准确率会从99%降到90%附近但这才是模型真实的泛化水平。5.2 Loss一直降准确率却卡在80%不动现象训练集Loss持续下降到0.1以下但验证集准确率在80%附近波动怎么调学习率都没用。原因CWRU数据集里并不是每个类别样本数量均衡。正常样本数量远多于滚动体故障样本模型只需要偏向预测正常样本就能把Loss压得很低但滚动体故障的召回率很差。准确率80%可能正是模型把所有样本都预测为正常的比例。这里用第4章的混淆矩阵指标就很容易发现。解决给损失函数加类别权重常用的做法是用torch.nn.CrossEntropyLoss(weightclass_weight)class_weight按样本数量的倒数计算或者把少数类做训练增强。改完以后整体准确率可能略有下降但漏检率会显著改善。我从那以后每次训练都会先统计一遍训练集的类别分布。5.3 TensorBoard启动后看不到任何曲线现象运行tensorboard --logdir logs --port 6006后浏览器里显示No dashboards are active或者找不到事件文件。原因一种情况是logs目录下还有子目录而--logdir指向的路径里没有直接包含.tfevents文件另一种情况是目录里确实是空的训练时没启用TensorBoard写入。这套源码的logs里有很多events.out.tfevents.*文件但分布在不同的实验子目录中。解决先把--logdir指向有事件文件的祖父目录或者用find logs -name *.tfevents确认事件文件位置。如果日志确实在TensorBoard是可以递归读取的你需要确认路径没有权限问题。如果子目录太深也可以临时把所有事件文件复制到logs_all目录下再启动。5.4 STFT时频图模糊一片看不出调制边带现象运行draw_transform.py保存的STFT图颜色糊成一团连冲击特征都分不清更别说轴承故障的调制边带了。原因STFT窗口长度和重叠率设置不合适。窗口太短频率分辨率低频带是糊的窗口太长时间分辨率低冲击位置看不出来。另外如果直接用矩形窗频谱泄漏会非常严重必须加汉宁窗。项目里明确提到STFT使用汉宁窗说明作者也踩过这个坑。解决我一般设置窗口长度为256重叠75%。窗口长度越长频率分辨率越高重叠越多时间平滑越好。对12kHz采样率的CWRU信号窗口256对应约21毫秒能保留较好的冲击细节。修改后重新生成时频图你会看到内圈故障的调制频率带清晰很多。这个参数可以直接在draw_transform.py里搜索hop_length和n_fft修改。6. 进阶用draw_transform.py验证特征可分性再回头调模型6.1 CWT和STFT一次可视化就能决定模型选型draw_transform.py对CWRU数据进行了两种时频变换CWT连续小波变换和STFT短时傅里叶变换汉宁窗并把结果画出来。这个脚本的价值不只是“好看”它可以直接帮你判断当前数据适合用CNN还是AE、应该用哪种预处理。我的使用习惯是拿到新数据后先跑一遍draw_transform.py对比正常样本和内圈故障样本的时频图。如果在CWT图上能看到明显的冲击间隔而STFT图上只看到频带抬升那么用AE做重构检测会更容易因为重构误差在冲击位置会被放大如果STFT图上有清晰的边频带说明调制信号是稳定的周期成分这时用CNN分类更靠谱。CWT和STFT各有侧重不要偷懒只画一种。6.2 把可视化结论转化为模型改进一个可复现的技巧具体怎么转化这里有一个很实用的做法。如果你发现正常样本和故障样本在STFT图的某些频带差异显著就可以把预处理从原始波形换成STFT时频图然后将二维时频图作为CNN的输入替代一维原始信号。常见做法是用torchvision.transforms把时频图缩放成224×224再输入ResNet等二维网络。以本项目为例你可以在CNN_Datasets里新增一种数据调用draw_transform.py生成好的STFT灰度图存储为numpy矩阵或PNG图片。训练时加载这些图而不是加载原始波形。这样CNN能同时学到时间维和频率维的特征对滚动体故障的早期微弱冲击更敏感。我做过对比同样的CNN结构输入从一维原始波形换成STFT图后漏检率能降低两到三个百分点。从那以后我每次拿到类似的轴承数据集都会强制自己先做一遍时频可视化再决定网络和预处理。这个习惯避免了我很多次盲目调参的无效劳动。希望这几条实战经验能帮到你也建议你把这份源码下载下来边跑边改踩几次坑之后你对深度学习故障检测的理解会比只看论文深得多。本文还有配套的精品资源点击获取