简介围绕CWRU轴承数据集这份PyTorch源码包提供了自编码器AE与卷积神经网络CNN等多种网络的实现、训练代码及项目说明面向深度学习故障检测方向的研究者与学习者便于研读网络结构、对比不同数据预处理方式并开展相关实验。资源共478个文件压缩包约1.16MB主体为254个Python脚本含模型定义、训练流程与可视化脚本另有tensorboard日志、配置文件等目录按数据预处理、模型定义、工具函数、训练日志等功能模块划分层次清晰。已有874人学习下载。代码在原始版本基础上做了改造增强训练脚本中新增TensorBoard可视化与精确率、召回率、误报率、漏检率等指标记录并补充了模型ACC/LOSS曲线绘制脚本、CWRU数据CWT和STFT变换分析脚本。整套资源既可帮助初学者通过PyTorch框架理解故障检测建模流程也能让有基础的读者快速获取可运行源码、日志与指标输出用于复现实验或改进算法。1. 多模型故障检测源码实测CWRU轴承数据、PyTorch训练链路与TensorBoard日志可视化把CWRU轴承数据集上的故障检测跑通比单纯跑通一个图像分类要多花不少心思。这份基于多种深度学习的故障检测算法python源码把CNN和自编码器两类网络、三种数据预处理方式、完整的检测指标与可视化脚本打包到一起训练日志、TensorBoard曲线、CWT与STFT时频变换图都齐了。它不解决端到端部署但很适合想研究故障检测算法、想把深度学习真正落到振动数据上的人。对还在折腾环境、刚接触PyTorch的新手来说这份代码也能当一套标准的训练流水线来模仿。读完后你能回答三个问题输入信号怎么切、网络怎么选、训练日志里哪些指标才真正说明检测能力。2. 拿到源码先拆目录六个模块职责、最小环境配置与数据组织2.1 目录结构与模块的调用关系解压之后先别急着跑train.py。我习惯先用tree把整个目录扫一遍弄清楚每个文件夹是谁在喂谁、谁的输出落在哪里。这份资源和很多散装demo不一样的地方是它把“数据预处理、网络定义、训练辅助函数、可视化脚本”分得比较清楚改一个环节不用牵动全部代码。顶层结构大致如下├── AE_Datasets/ # 自编码器输入三种数据预处理方式的相关代码 ├── CNN_Datasets/ # CNN输入三种数据预处理方式的相关代码 ├── checkpoint/ # 各网络训练过程保存的日志作者跑实验时的留痕 ├── logs/ # TensorBoard事件文件events.out.tfevents.* ├── models/ # 不同网络模型定义CNN、AE等 ├── utils/ # 训练过程辅助函数train_utils.py、train_utils_ae.py ├── draw_models.py # 各模型训练集/验证集ACC与LOSS绘图 ├── draw_transform.py # CWRU数据的CWT/STFT变换分析与绘图 ├── train.py # 训练除自编码器以外的网络模型 └── train_ae.py # 训练自编码器网络模型从调用关系上看AE_Datasets和CNN_Datasets是数据入口分别服务自编码器和CNN。models里的网络结构负责接收数据并输出预测utils里的train_utils.py和train_utils_ae.py封装训练逻辑train.py和train_ae.py是最外层入口checkpoint和logs是训练过程中写到磁盘的产物两个draw脚本把日志和原始数据变成可视化的图。也就是说如果你只想看某个网络在CWRU上跑到什么水平只需要改train.py里模型实例化那一行如果你想对比不同预处理方式先在两个Datasets目录里把数据准备好再启动训练。checkpoint里的日志是作者当时跑实验留下的我建议先留着不要删后面跟自己的实验对照时很有用。2.2 最小可跑环境PyTorch、TensorBoard与依赖安装代码基于PyTorch训练过程用了TensorBoard写入Scalar指标可视化用到matplotlib和numpy。故障检测常用的时频变换依赖scipy和pywt后者不是PyTorch自带的跑draw_transform.py之前必须装。我一般用conda建一个独立环境避免把系统Python搞得一团糟conda create -n fault_detect python3.8 conda activate fault_detect pip install torch tensorboard numpy scipy matplotlib pip install pywt第一行创建Python 3.8环境第二行激活环境。PyTorch的安装建议根据自己的CUDA版本选择CPU机器直接装CPU版也能跑只是速度慢一些torch版本不要求最新TensorBoard的events.out.tfevents.*格式在各版本间是兼容的。pywt是时频分析必须的一环装完可以用下面这条命令确认导入正常python -c import torch, tensorboard, scipy, pywt; print(ok)只要不报ImportError环境就算搭好了。这里有个容易被忽略的点如果电脑上同时装了多个Pythonconda环境激活后还要确认pip对应的就是当前环境否则很容易出现包装到了base环境、训练时却报ModuleNotFoundError的情况。Windows笔记本上跑这份代码没有额外障碍logs目录里那串LAPTOP-1FVELO7I主机名说明作者当时就是在笔记本上完成的实验。2.3 三种数据预处理方式同一种网络在不同输入下的效果差异AE_Datasets和CNN_Datasets都写了三种数据预处理方式这是这份资源里比较有价值的对比实验点。常见的做法是把CWRU原始振动信号按固定窗口长度切段三种方式的差别主要在于切法和归一化上我按最常见的一组配置给一个对照。预处理方式输入形态典型适用网络特点原始振动段一维时序数组AE、1D-CNN、LSTM信息最直接训练快对噪声敏感归一化/标准化振动段一维时序数组AE、1D-CNN、LSTM保留波形形态减小幅值差异影响时频变换谱图二维图像2D-CNN能看到时间与频率维度的能量分布计算量大这三种方式在CWRU上都可以跑但结论可能有差别。我的经验是一维输入训练快适合先验证网络结构和调参想追求更稳的检测效果可以在确认一维基线之后再做STFT或CWT谱图喂给2D-CNNdraw_transform.py就是专门用来提前观察信号变换效果的。这里也想提醒一句无论选哪种方式训练集和验证集都要在同一预处理管线里生成不能训练用归一化、验证用原始数据否则指标再好看都是假的。3. 训练入口拆解train.py与train_ae.py的工作流、检测指标与日志写入3.1 训练脚本骨架数据加载、模型实例化与epoch循环train.py和train_ae.py是两份独立入口逻辑高度相似。它们本身不含太多代码真正的训练函数在utils/train_utils.py和utils/train_utils_ae.py里脚本负责解析数据集、构建模型、把参数传进去。打开train.py核心流程大致是这样# train.py 的精简骨架完整逻辑以项目内源码为准 import torch from torch.utils.data import DataLoader from utils.train_utils import train # 训练封装函数 from models import build_model # 以models目录里实际提供的函数为准 # 数据加载CNN_Datasets里预处理好的数据通过Dataset包装 train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse) # 按经典CWRU四分类正常/内圈/外圈/滚动体构建网络 model build_model(cnn, num_classes4) train(modelmodel, train_loadertrain_loader, val_loaderval_loader, epochs100, lr1e-3, log_dirlogs/cnn_result, ckpt_dircheckpoint/cnn_result)这段代码的核心逻辑是把预处理好的数据包装成DataLoader训练时shuffleTrue打乱顺序验证时不打乱。模型通过models目录里的构建函数创建train.py负责传参train函数内部完成前向、反向、指标计算和日志写入。参数方面batch_size设64在大多数笔记本上跑得动数据量大再考虑128epochs设100是常见设置先跑50轮看曲线趋势也行lr1e-3是Adam类优化器的常见起点如果loss震荡明显就降到1e-4。train_ae.py的运行方式几乎一样区别只在模型换成了自编码器以及内部走train_utils_ae.py的封装自动编码器对重构误差的监控和对分类指标不太一样。3.2 故障检测核心指标精确率、召回率、误报率、漏检率的计算方式很多新手训练时只看准确率这在故障检测里是不够的。CWRU数据集中正常样本和故障样本数量往往不均衡如果故障样本只占少数模型全部预测正常也能拿到很高的准确率但这样的模型毫无检测能力。所以这份源码在训练函数里特意加了精确率、召回率、误报率、漏检率和F1值计算逻辑是二分类视角正类为故障负类为正常。# 把多分类预测映射成二分类0正常1故障 binary_true (y_true ! 0).astype(int) binary_pred (y_pred ! 0).astype(int) # 计算混淆矩阵四要素 TP ((binary_true 1) (binary_pred 1)).sum() FP ((binary_true 0) (binary_pred 1)).sum() FN ((binary_true 1) (binary_pred 0)).sum() TN ((binary_true 0) (binary_pred 0)).sum() precision TP / (TP FP) # 精确率报故障里真的故障比例 recall TP / (TP FN) # 召回率真故障里被抓住的比例 false_alarm FP / (FP TN) # 误报率正常样本被错判成故障 miss_rate FN / (FN TP) # 漏检率故障样本被漏掉 f1 2 * precision * recall / (precision recall 1e-8)指标含义很直观精确率回答“报出来的故障有多少是准的”召回率回答“真实故障有多少被抓住了”误报率关注正常样本被误伤漏检率关注故障样本被漏过。产线场景里误报率高了会导致频繁停机检查代价很大学术对比则更常用F1来平衡精确率和召回率。分母上加1e-8是为了防止除零这是数值稳定性的常规写法。训练结束后utils会把这几个数连同Loss值一起写入TensorBoard。你不需要自己算但必须明白一个道理acc高不代表检测能力强只有当precision、recall和F1同时稳定这个模型才真的可用。3.3 TensorBoard日志写入看Scalar曲线而不是等训练结束训练过程中train函数会向logs目录写入Scalar数据格式就是那串events.out.tfevents开头的文件。查看方式很简单tensorboard --logdirlogs启动后浏览器访问配置文件里提示的端口默认是6006。左侧Scalar面板能看到train_loss、val_loss、train_acc、val_acc以及精确率、召回率、误报率、漏检率、F1值这些曲线。比等100轮跑完再拿终端输出判断要直观得多loss下降平缓还是震荡val_acc到哪个epoch开始不再上升一眼就能看出来。logs目录里散落着多组events.out.tfevents.*文件文件名中的时间戳是Unix格式的启动时刻主机名后面是进程ID。多组文件意味着作者跑过多轮实验这其实是给后来者留下的一个好处你可以对比不同实验的曲线找出哪一组预处理方式或哪个网络参数组合的val曲线最稳。训练到一半发现val_loss反弹直接停掉改参数不必等全部跑完。4. 可视化落地draw_models.py画曲线、draw_transform.py做时频变换4.1 draw_models.py训练集与验证集的ACC/LOSS曲线绘制训练结束后终端只能看到最后一轮的数字中间过程全在日志里。draw_models.py做的就是把这些日志解析出来画成曲线图方便整体观察趋势。它背后的读取逻辑本质上就是解析TensorBoard事件文件核心思路是这样# draw_models.py 的读取逻辑绘图细节以项目源码为准 import os import matplotlib.pyplot as plt from tensorboard.backend.event_processing.event_accumulator import EventAccumulator for run in os.listdir(logs): ea EventAccumulator(os.path.join(logs, run)) ea.Reload() # 先看有哪些tag避免写错名字 tags ea.Tags()[scalars] train_acc [s.value for s in ea.Scalars(train_acc)] val_acc [s.value for s in ea.Scalars(val_acc)] plt.plot(train_acc, labelf{run}-train) plt.plot(val_acc, labelf{run}-val) plt.xlabel(epoch) plt.ylabel(acc) plt.legend() plt.savefig(acc_curves.png)EventAccumulator是TensorBoard官方提供的日志解析接口Reload方法把事件文件读进内存Scalars方法按tag取出一组标量。如果你的训练函数里tag命名不是train_acc先通过ea.Tags()[scalars]打印所有键名再改动。这个脚本的价值在于多个实验跑完后所有acc曲线叠加在一张图上哪个模型在第几个epoch过拟合、哪个曲线整体更平稳一目了然。LAOSS曲线同理只是把tag换成train_loss和val_loss。我通常把ACC和LOSS两张大图都导出来放一起看ACC稳定但LOSS持续走高的组合很值得警惕。4.2 draw_transform.pyCWT与STFT汉宁窗分析CWRU振动信号CWRU数据集里的原始数据是轴承振动加速度信号采样频率有12kHz和48kHz两种。直接看时域波形内圈故障和外圈故障的冲击差异并不明显把信号变换到时间-频率平面后不同故障的特征频率带会清晰许多。draw_transform.py做的是这件事它用短时傅里叶变换和小波变换比较同一段信号的两种时频表示。# draw_transform.py 核心变换逻辑以CWRU 12kHz采样为例 import numpy as np from scipy.signal import stft import pywt fs 12000 # 采样频率48kHz数据请同步修改 # 方式一STFT汉宁窗窗口256点重叠128点 f, t, Zxx stft(signal, fsfs, windowhann, nperseg256, noverlap128) plt.pcolormesh(t, f, np.abs(Zxx), shadinggouraud) # 方式二CWTmorlet小波尺度范围1~64 scales np.arange(1, 64) coefs, freqs pywt.cwt(signal, scales, morl, sampling_period1/fs) plt.imshow(np.abs(coefs), aspectauto, cmapjet, extent[0, len(signal)/fs, freqs[-1], freqs[0]])STFT是把长信号切成一段段短窗口做傅里叶变换汉宁窗能抑制频谱泄漏窗口长度决定时间分辨率和频率分辨率的折中窗口越长频率越精细但时间定位越模糊。CWT用morlet小波在不同尺度上做卷积低频尺度对应频率高、时间分辨率差高频尺度则相反整体上比STFT更擅长捕捉瞬态冲击。变换窗口/小波输出维度主要看什么STFT汉宁窗256点窗口重叠128点频率×时间稳态成分和冲击的频段范围CWTmorlet尺度1~64尺度×时间瞬态冲击的精确定位在CWRU场景里轴承故障会周期性产生冲击这些冲击在时频图上会形成沿时间轴重复的能量线条特征频率和故障类型一一对应。我第一次把内圈故障样本跑完STFT后看到低频段有规律的亮带对照故障特征频率公式一算完全对得上那种感觉比直接看测试集acc踏实得多。4.3 checkpoint、logs和events时间戳一次训练的三种留痕checkpoint和logs经常被混为一谈其实分工不同。checkpoint目录是训练过程中控制台输出的日志文本记录了每个epoch的Loss、acc、precision、recall等数字适合快速翻看logs目录是TensorBoard事件文件适合可视化对比。两者时间戳一致时对应同一次训练。logs目录里这组events.out.tfevents.*文件的时间戳集中在1671780001到1671783488之间按照Unix时间换算不到一个小时说明作者在连续多轮实验后把结果留了下来。我习惯用这个时间戳做实验编号每次训练在log_dir里带上日期和模型名比如logs/20250214_cnn_stft配合checkpoint文本日志三个月后回来看还能还原当时的调试过程。5. 训练避坑日志对不上、曲线空白、时频图发黑这类问题的排查记录5.1 现象TensorBoard里曲线空白Scalar面板没有数据第一次打开tensorboard时Scalar面板空白是常见翻车现场。最常见的原因是logdir传错了层级logs目录下直接放事件文件但如果没有按实验分子目录TensorBoard扫描到的是同一层级下的多个事件文件可能只显示最后一段。另一个原因是训练还没跑够一个epochwriter没有flush事件没有真正落盘。解决方法是训练脚本确认每个实验写入独立的子目录比如logs/cnn_run1启动tensorboard时用tensorboard --logdirlogs如果跑了好几个epoch还空白检查代码里writer.close()或flush()是否执行。事件文件是追加写入的训练中途强杀进程也可能导致最后一小段数据缺失但前面的epoch应该还在。5.2 现象训练集acc很高验证集acc明显偏低这个现象十有八九是数据切分出了问题而不是网络过拟合。CWRU数据来自同一个轴承的连续振动采样如果随机把窗口切进训练集和验证集同一个原始文件里的相邻窗口可能同时出现在两侧本质是数据泄漏。训练时模型见过这些窗口的邻居验证自然得分高换个工况就被打回原形。解决方法是按样本文件切分而不是按窗口切分先按CWRU原始数据文件划分train/val再从文件内部切窗口。我自己踩过这个坑当时训练集acc到了97%以上验证集只有83%改成按文件切分后两边的差距明显收窄。特征频率的计算也要以同一采样率为前提12kHz和48kHz的数据混着用会直接破坏一致性。5.3 现象CWT/STFT图一片黑或者亮得刺眼先确认画的是np.abs(系数)还是原始复数很多新手直接pcolormesh复数数组matplotlib会报警告或者画出诡异色块。排除这个之后一片黑通常是幅值动态范围问题信号里有低频分量幅值很大把其他频段压成了黑色。STFT对这种问题相对宽容CWT尤其明显因为小波系数的幅度随尺度变化很大。解决方法是做对数压缩或者线性归一化。我在处理CWRU数据时习惯对系数取20*log10后的值再画图动态范围会舒服很多也可以把cmap的vmin/vmax手动设到百分位数比如vmin0.5分位vmax99.5分位这样能避开个别极大值把色标拉伸过头。这类问题多半不是数据坏了而是显示方式没调对。5.4 现象准确率很高但误报率和漏检率居高不下这个现象很多人会愣一下因为它恰好是类别不平衡的典型特征。CWRU里如果正常样本占比大模型学到的决策边界会偏向多数类acc被正常样本撑高了但故障样本漏检严重误报率也可能异常。单纯盯着acc调参永远走不出来。解决方法是把评估重心转移到precision、recall、F1和混淆矩阵上训练时也可以考虑给少数类加权。这份源码在utils里加了误报率和漏检率的计算就是明确暗示故障检测任务不能只看acc。如果你自己训练的模型出现这种状况先去看混淆矩阵然后决定是加权还是换输入形式。5.5 现象CPU训练很慢数据加载占了绝大部分时间没有GPU时训练慢是正常的但有一种慢完全可以避免训练循环里每个step都做STFT或CWT变换相当于把预处理计算重复做了几十遍。CWRU数据总量不大但如果每个窗口都现场做时频变换CPU训练会慢得让人怀疑人生。解决方法是把预处理结果提前算完存成npy训练时直接加载数组。第一次跑draw_transform.py确定用哪种变换后我习惯先把整批数据变换好写入磁盘再让训练脚本读取训练速度能提升一个量级。另一个常见卡点是DataLoader的num_workers在Windows上设置不当导致黑屏或死锁Windows建议设0或1。6. 用事件日志做验证解析logs批量对比指标固定数据入口再谈调参6.1 批量解析logs目录用事件文件生成指标对比表训练脚本每跑一次就往logs写一批事件文件时间一长手动翻找某个实验的指标很麻烦。我常用的一个技巧是写一段解析脚本把logs下所有实验的最后一个epoch指标抽出来自动生成对比表省去反复开关TensorBoard的功夫。# 批量提取logs下各实验的最终指标tag名以实际写入为准 import os from tensorboard.backend.event_processing.event_accumulator import EventAccumulator for run in sorted(os.listdir(logs)): ea EventAccumulator(os.path.join(logs, run)) ea.Reload() val_acc [s.value for s in ea.Scalars(val_acc)] val_f1 [s.value for s in ea.Scalars(val_f1)] print(f{run:24s} val_acc{val_acc[-1]:.4f} val_f1{val_f1[-1]:.4f})这段代码的价值在于把对比从“打开TensorBoard手动记”变成“一条命令出表格”。每一行代表一次实验val_acc和val_f1取最后一个epoch的值。如果想要更完整的结论把误报率和漏检率也加进去。多组实验一次性摊开哪个网络、哪种预处理方式表现稳立刻有数不用凭记忆下结论。6.2 把曲线和时频图回放成调参依据固定数据入口再谈微调调参数有一个很容易犯的错一边改预处理、一边改学习率最后指标变了也说不清是哪个改动起的作用。我的做法是先用draw_transform.py把CWRU样本的STFT和CWT图看一遍确定用哪种输入形式然后把这一项固定下来之后所有对比都只动网络结构和训练参数。选择网络时也不只看最终acc把draw_models.py画出来的多组曲线叠加在一起验证集acc在第几个epoch见顶、之后是否下滑比单个数字更有信息量。如果某个模型在第30轮就开始过拟合即使最终acc看着不错几轮后换工况也很可能翻车。从那以后我再拿到一份故障检测源码都不会直接train.py一把梭而是先把logs目录里的旧实验全部解析出来画成对比表结合时频变换图确认数据入口再决定要不要动手改参数。这个习惯让我少走了不少弯路。希望帮到你。本文还有配套的精品资源点击获取