简介:本资源为基于深度学习的滚动轴承故障诊断Python项目,面向机械、自动化及计算机方向的高年级本科生和研究生,适用于毕业设计、课程作业或故障诊断入门实战。项目围绕CWRU轴承数据集展开,提供数据预处理、特征提取、模型训练与可视化分析等完整流程,难度适中,便于理解深度学习在工业信号处理中的应用。压缩包共41个文件,以mat数据文件、py脚本和md说明文档为主,约34.87MB,其中Python脚本覆盖数据读取、模型搭建和散点图绘制等环节。目前已有74人学习下载。项目源码均经过本地编译调试,评审得分98分,能稳定运行;随包附带目录说明与README文档,可帮助使用者快速定位关键模块、复现实验流程。
1. 基于深度学习的滚动轴承故障诊断:这份 Python 源码包能帮你把毕业设计从“能跑”做到“能答辩”
做滚动轴承故障诊断的毕业设计,最容易遇到的尴尬是:信号处理那一套(时域统计量、FFT 频谱、包络谱)写了两万字,最后模型却只有一个 SVM 或随机森林,答辩时被问“为什么不用深度学习”直接卡壳。换深度学习方案又怕数据量不够、训练不稳、调参玄学。这份基于深度学习与 Python 的滚动轴承故障诊断源码包,走的是更务实的一条路:把振动信号转成时频谱图,用预训练 ResNet 做迁移学习,从数据预处理、模型训练到评估可视化给出了完整闭环,而且附带了训练好的模型权重,适合拿来当毕业设计主干,也适合刚接触故障诊断 + 深度学习的同学做基线系统。我拆完这套源码后最大的感受是:它的分工很清晰,哪部分是数据脚本、哪部分是模型定义、哪部分是训练入口,一目了然,照着跑一遍就能摸清整个流程的骨架。
2. 数据准备与预处理:从原始振动信号到模型输入
2.1 数据集组成与标签体系:先搞清你手里有什么
这套源码用的是某公开轴承数据集,驱动的加速度计信号,采样频率 12 kHz,包含了四种健康状态:正常、内圈故障、外圈故障、滚动体故障。每种故障又按损伤直径分为 0.007、0.014、0.021 英寸三档,加上不同负载工况(0 到 3 hp),组合下来有 10 个左右的类别。这是做诊断任务最常见的数据组织方式,也是毕业设计里最容易讲清楚的一个点。
源码里把原始数据整理成了如下结构:
dataset/ ├── 0_Normal/ │ ├── normal_0_0.csv │ └── normal_0_1.csv ├── 1_IR_0.007/ │ ├── ir_0.007_0.csv │ └── ir_0.007_1.csv ├── 2_OR_0.014/ ├── 3_BF_0.021/ └── ...每个 CSV 文件就是一段连续的振动信号采样点,一列数值,没有额外通道。文件名里的第一个数字是类别编号,第二个是同一工况下的重复采集批次。
这里有一个关键点:类别编号直接对应训练时的 label,所以第一步是检查数据目录里的类别数量和编号连续性。如果某个故障尺寸的样本缺失,训练脚本会报 one-hot 维度不匹配,这类问题在跑源码时很容易遇到。我一般会先写一个几行的脚本把所有类别遍历一遍,打印出每个文件夹的文件数和类别 ID,再决定要不要手动补标签。
import os data_root = "dataset" classes = sorted(os.listdir(data_root)) for idx, cls in enumerate(classes): file_list = os.listdir(os.path.join(data_root, cls)) print(f"类别ID: {idx} | 故障类型: {cls} | 样本文件数: {len(file_list)}")逻辑说明:这段代码做的事情是把类别文件夹按名称排序,按顺序分配整数 ID,检查每个类别的样本数量是否均衡。实际运行时如果发现某一类文件数特别少(比如内圈 0.021 英寸只有几段),后续训练就要考虑数据增强或者类别加权,否则模型会天然偏向样本多的类别。
有了目录和标签之后,接下来的步骤是把一维振动信号切片,变成模型能吃的一批批数据。
2.2 滑窗切分与训练/测试集划分:防数据泄漏的第一道防线
原始振动信号很长,不能整段直接塞进模型,常见做法是用固定长度的滑窗去切。这套源码里滑动窗口长度设的是 2048 个采样点,滑步(stride)1024,也就是说相邻两个窗口有 50% 的重叠。这个设置的好处是:每个样本的长度足够覆盖多个旋转周期,而重叠能增加样本量,对数据不充裕的故障类别比较友好。
import numpy as np def sliding_window_slice(signal, window_len=2048, stride=1024): samples = [] start = 0 while start + window_len <= len(signal): samples.append(signal[start:start + window_len]) start += stride return np.array(samples)参数说明:窗长 2048 在 12 kHz 采样率下对应约 0.17 秒的物理时长,够包含多个轴承旋转周期;stride 1024 是窗长一半,重叠率 50%,样本量大约翻倍。如果想做消融实验,这两个参数是最值得调的——窗太短会丢失低频故障特征,窗太长又会引入过多噪声。
切分完成之后就是划分训练集和测试集。这一步是整套流程最容易出错的地方,也是毕设答辩老师最爱抓的漏洞:按样本随机划分会导致数据泄漏。
所谓数据泄漏,就是同一个原始信号切出来的相邻窗口可能同时出现在训练集和测试集里,模型相当于“见过”测试数据的一部分。这样一来测试准确率虚高,换到现场数据立刻打回原形。正确做法是按“原始记录文件”划分,也就是同一个 CSV 切片出来的所有窗口必须待在同一个集合里。
from sklearn.model_selection import train_test_split file_names = [f for f in os.listdir(data_root + "/1_IR_0.007") if f.endswith(".csv")] train_files, test_files = train_test_split(file_names, test_size=0.2, random_state=42)逻辑说明:这里把文件名列表按 8:2 拆开,然后用文件名去逐个读入信号并切片。这样保证测试集里的窗口全部来自训练时没见过的原始记录文件,评估结果才可信。random_state 固定成 42,复现实验时保持一致。
2.3 信号转时频谱图:STFT 参数怎么定、为什么不用原始波形
模型输入不是一维数组,而是二维的时频谱图。原因很简单:一维卷积对振动信号的特征提取能力有限,而预训练图像模型(ResNet 系列)是在大量自然图片上训练出来的,直接吃时频谱图能借到很强的二维特征提取能力。这也是本源码方案的核心思路——用短时傅里叶变换(STFT)把一维信号变成“图片”。
STFT 的核心参数是窗函数、窗长和重叠。源码里用的是 Hanning 窗,窗长 256,hop length 128。
import librosa def signal_to_spectrogram(signal, n_fft=256, hop_length=128, sr=12000): stft = librosa.stft(signal, n_fft=n_fft, hop_length=hop_length, window="hann") spectrogram = np.abs(stft) # 转成对数幅度,压缩动态范围 log_spectrogram = librosa.amplitude_to_db(spectrogram, ref=np.max) return log_spectrogram参数说明:n_fft=256 就是做 FFT 的点数,决定了频率分辨率,12 kHz 采样率下每个频率 bin 约 46.875 Hz。这个分辨率对轴承故障特征频率(通常在几十到几百 Hz 的范围)是够用的。hop_length=128 是帧移,控制时频谱图的时间分辨率,和 n_fft 的比例决定了相邻帧的重叠率。输出的是 dB 刻度的对数谱,原因是原始幅度谱动态范围太大(可能跨两三个数量级),直接送进网络会让模型过分关注幅值大频率带。
转成对数幅度谱之后,还要做一步归一化,把数值范围压到 0 到 1 之间。因为 ResNet 的预训练权重是基于 ImageNet 归一化统计的,虽然迁移学习时网络会自适应调整,但输入分布差异过大会拖慢收敛速度。源码里在 DataLoader 里做了标准化,均值方差用的是 ImageNet 的默认值,这是一个值得注意的细节。
3. 模型选型与训练:ResNet 迁移学习为主干方案
3.1 为什么选 ResNet:从一维 CNN 的翻车说起
我第一次做这个任务时图省事,自己搭了一个四层一维卷积网络,结构类似 LeNet 的一维版本,输入 2048 点原始信号,最后接全连接层做分类。结果很惨:训练集准确率很快就到了 99%,但验证集只有 87%,过拟合非常严重。原因也不复杂,故障振动信号的样本量本身不大(滑窗后也就几千个样本),一维 CNN 的可学习参数太多,而原始信号里的特征又非常微弱,模型很容易记住噪声模式。
换成预训练 ResNet18 之后,情况完全变了。预训练模型的特征提取器已经能在二维图像上提取纹理、边缘类特征,时频谱图上的故障特征(周期性冲击在频域的表现是等间距谱线)恰好属于这类模式。迁移学习让模型从“从零学习”变成了“在已有视觉特征基础上微调”,收敛速度和最终精度都明显提升。
3.2 模型结构:ResNet 输入怎么适配单通道时频谱图
ResNet18 默认接收 224×224×3 的 RGB 图像,而时频谱图是单通道灰度图。源码的处理方式是:将时频谱图复制三次变成三通道。这是一个简单但有效的做法,比修改第一层卷积的输入通道更省事,也方便后续换其他预训练模型时保持一致。
import torchvision.models as models import torch.nn as nn def build_resnet_model(num_classes=10): model = models.resnet18(pretrained=True) # 替换最后一层全连接,输出类别数 in_features = model.fc.in_features model.fc = nn.Linear(in_features, num_classes) return model逻辑说明:把最后一层全连接从 1000 类换成 10 类。in_features 是从原模型读出来的,不硬编码为 512,这样以后换成 ResNet34 或 ResNet50 时不需要改这行代码。
训练时还有一个微调策略的问题。两种常见做法:一是冻结全部卷积层只训练最后的全连接层,适合数据量很小的场景;二是只冻结前面几层,后面几层跟着微调。这套源码默认采用第二种策略——前两层残差块冻结,后面的层保持可训练。原因很简单:时频谱图的纹理特征和自然图像的低层特征(边缘、角点)是共通的,但高层语义差异很大,所以高层需要重新适应。
for name, param in model.named_parameters(): if name.startswith("layer1"): param.requires_grad = False参数说明:这行代码把 layer1 之前的参数全部冻结,因为 ResNet 的命名规则里 stem(第一层卷积)和 layer1 是连续的。实际训练时建议打印一下各层 requires_grad 是否和你预期一致,防止误冻结了关键层。
3.3 训练参数详解:优化器、学习率、Batch Size 怎么配合
训练参数直接决定了模型能不能收敛、会不会震荡。源码里选的是 AdamW 优化器,初始学习率 1e-4,weight decay 设为 1e-4,batch size 32,训练轮数 50。
import torch from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4) scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6) for epoch in range(50): model.train() for inputs, labels in train_loader: optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() scheduler.step()逻辑说明:AdamW 相比原始 Adam 最主要的变化是 weight decay 不作用于动量项,这对迁移学习场景更友好,泛化能力略好。1e-4 的学习率是微调预训练模型的安全起点,经验上从 1e-3 起步很容易把预训练权重冲坏。weight decay 1e-4 是为了抑制过拟合,因为训练样本总量不大,不加正则的话全连接层很容易记死。CosineAnnealingLR 让学习率从 1e-4 按余弦曲线一路下降到 1e-6,前半程保持一个相对高的学习率让模型快速适应故障数据,后半程低学习率做精细收敛。
这里有一个容易被忽略的参数:batch size 选 32 除了显存考虑之外,还和数据集的类别数有关。如果某个类别的样本量少于 batch size,一个 batch 里可能只覆盖到部分类别,BN 层的统计量会不稳定,导致训练震荡。我一般会在训练前检查每个类别的最小样本数,如果小于 batch size 的 1.5 倍,就考虑做类别重采样或者缩小 batch size。
训练过程中的监控也很关键。源码里每个 epoch 结束都会在验证集上计算准确率,并保存验证集准确率最高的权重。这一点看起来基础但很多初写训练脚本的人会漏掉,最常见的就是训练完之后再用最后一轮权重去做测试,结果发现测试分数低于历史最优,白白浪费了模型。
4. 评估与可视化:准确率之外还要看什么
4.1 从准确率到混淆矩阵:分类器的体检报告
毕业设计答辩时,只亮一个“测试集准确率 98.7%”其实是很危险的。老师的下一个问题往往是:“你分类器的错误集中在哪些类别之间?”没有混淆矩阵的话这个问题很难回答。源码里专门写了评估脚本,输出每类的精确率、召回率、F1 值和混淆矩阵。
多分类问题里,准确率不是唯一指标,甚至不是最重要的指标。故障诊断场景更关注漏报和误报。漏报的代价是故障轴承继续运行直到彻底损坏,误报的代价是停机检查浪费时间。所以做评估时我习惯看两个指标:宏平均 F1 和加权平均 F1。宏平均不考虑类别样本量,能反映出小样本类别的真实表现;加权平均靠样本量加权,反映总体性能。
4.2 训练曲线与混淆矩阵的配合解读
训练曲线的解读是有一些经验的。最常见的病态曲线是:训练 loss 一路下降到很小,验证 loss 先降后升,形成“开口向上”的形态。这就是过拟合的典型信号,在学习率调得比较大的时候尤其明显。
混淆矩阵的输出格式如下:
预测为正常 预测为内圈 预测为外圈 预测为滚动体 实际为正常 120 0 0 0 实际为内圈 0 116 4 2 实际为外圈 1 12 90 5 实际为滚动体 0 3 6 110从这个矩阵能读出很多信息。内圈故障和外圈故障的混淆最严重——因为它们都是滚道缺陷,特征频率接近,在低负载工况下冲击特征差异不明显。滚动体故障的召回率最低(漏报最多),原因是滚动体自转周期长,缺陷进入承载区的时间占比低,信号中的周期性冲击更微弱。这些结论写在毕设论文里,比一句“准确率高”有说服力得多。
4.3 特征可视化:t-SNE 检验特征可分性
还可以把模型倒数第二层的特征向量拿出来做 t-SNE 降维,看看不同类别在二维平面上的分布。这一步能直观展示模型是否真的学到了可分特征,也是答辩时很出彩的一张图。
from sklearn.manifold import TSNE import matplotlib.pyplot as plt features = extract_features(model, dataloader) # shape: (N, 512) tsne = TSNE(n_components=2, perplexity=30, random_state=42) features_2d = tsne.fit_transform(features) plt.scatter(features_2d[:, 0], features_2d[:, 1], c=labels, cmap="tab10", s=10) plt.show()逻辑说明:extract_features 是把模型的 fc 层临时替换成恒等映射,前向传播时直接取出池化层的输出。perplexity=30 是 t-SNE 最常用的取值,样本量小于 1000 时可以考虑降到 20 或 15,避免簇结构过散。
t-SNE 图里如果同类样本聚成几个独立小簇,而不是一个紧凑簇,说明特征里还混合了其他变量,比如负载工况的信号强度差异。这不是错误,但需要在答辩时主动解释:负载变化导致振动幅值整体偏移,特征空间里同类不同负载的样本自然分开,这是合理的领域差异,不是模型缺陷。
5. 避坑记录:数据泄漏、样本标签错位与迁移学习尺寸问题
5.1 坑一:按样本随机划分数据集导致测试准确率虚高
现象:测试集准确率超过 99.5%,换到另一批同型号轴承数据测试,直接掉到 82%。
原因:切分训练/测试集时直接随机打乱了所有滑窗样本。同一个原始信号切出的相邻窗口有 50% 重叠,模型见过测试窗口“一半的内容”,等于是开卷考试。
解决:改成按原始记录文件切分。写一个文件级 split 函数,把文件名列表按比例拆成训练和测试两份,再分别读入切片。划分完以后打印检查一下:训练集和测试集的窗口来源文件 ID 不得有交集。
5.2 坑二:STFT 之后忘记做对数压缩,模型收敛慢
现象:loss 下降很慢,前 10 个 epoch 准确率一直徘徊在 30% 左右。
原因:直接把线性幅度谱送进了模型。原始振动信号的频谱幅度动态范围很大,低频分量可能比其他频带高出两个数量级,模型被迫去拟合一个极不均衡的数值分布,BN 层的统计量也很难稳定下来。
解决:用 librosa.amplitude_to_db 把幅度谱转成 dB 刻度,动态范围从线性压缩到对数空间。变换后再用均值减法做标准化,数值范围基本落在 -50 到 0 之间,模型收敛速度明显加快。
5.3 坑三:ResNet18 的 224×224 输入尺寸和频谱图尺寸不匹配
现象:模型跑起来了,但验证准确率始终上不去,卡在 80% 上下。
原因:代码直接调用 ResNet 默认的 224×224 输入,而 STFT 输出的频谱图是 257×128 左右(取决于 n_fft 和信号长度)。盲目的 resize 把频谱图在一个方向上做了非等比拉伸,频率轴被压缩,高频细节糊成一片。
解决:在 STFT 参数里直接调整输出尺寸,让时频谱图的长宽比接近方形再做 resize。更省事的做法是把频谱图先 pad 再 resize,保留原始频率分量的比例。我一般直接调整 hop_length 让时间轴长度约等于 224,再配合等比 resize。
5.4 坑四:AdamW 的 weight decay 对预训练特征层影响过大
现象:训练前期正常,15 个 epoch 之后验证 loss 开始缓慢上升,但训练 loss 还在下降。
原因:对全部参数一刀切地加 weight decay,预训练层的权重被正则项压得太狠,之前学到的视觉特征被慢慢“遗忘”。全连接层的参数需要正则,但预训练卷积层的权重本身已经被规范化过了,不适合再施加同样的强度。
解决:把参数按 requires_grad 分组,冻结层直接用预训练权重不做正则,微调层和全连接层才施加 weight decay。代码实现上就是两次 add_param_group,然后在优化器初始化时区分对待。
5.5 坑五:数据增强过度,改变物理含义
现象:训练时加了随机平移和噪声扰动,结果验证准确率不升反降。
原因:振动信号是有物理含义的,某些时域增强操作会破坏故障冲击的周期性。比如随机时间平移会让同一个样本的相位关系错乱,模型无法稳定地学到冲击间隔特征。
解决:给这场景的数据增强做减法。可以保留的只有幅值缩放(模拟传感器灵敏度差异),但缩放系数范围要约束在 0.9 到 1.1,不能过猛。移位、翻转、添加高斯噪声这类增强方法需要逐一验证效果,不能直接照搬图像分类的经验。
6. 模型导出与推理:把训练好的故障分类器实际用起来
拿到了验证集最高准确率的权重,只做评估还不够,得把它导出成可以部署的格式。这一步在毕设答辩时很加分。TorchScript 和 ONNX 是两种最常用的导出格式。TorchScript 的好处是保持 PyTorch 生态的完整性,后续想继续微调也方便;ONNX 的好处是能转到其他推理框架。
model.eval() dummy_input = torch.randn(1, 3, 224, 224) traced_script_module = torch.jit.trace(model, dummy_input) traced_script_module.save("bearing_resnet18.pt")逻辑说明:trace 模式需要给一个固定尺寸的 dummy input,用实际输入张量跑一遍前向,把整个计算图记录下来。这里需要注意 batch 维度必须写 1,后续推理时输入的 batch 大小得和 trace 时保持一致,否则会报维度不匹配。
推理端的滑窗设计和训练时保持一样的参数,这是很多人容易忘的点。如果训练时窗长是 2048、stride 是 1024,推理时也必须一样,不然后台处理完的频谱图和模型见过的分布不一致,精度会莫名其妙掉下来。
推理时另一个值得做的改进是置信度阈值。Softmax 输出的概率分布并不总是“非黑即白”,某个样本可能对两个类别的概率都在 40% 左右。这种情况说明模型认为样本特征不明确,直接取最大值判类别容易误判。我习惯设置一个 0.7 的阈值,低于这个值的输出归为“未知/待人工复核”。这个逻辑在实际使用场景里很重要——现场不只有训练时见过的那十类状态,还可能出现不对中、松动、润滑不良等未训练过的故障模式,模型应该能识别出“我不知道”,而不是硬套一个类别。
从那以后我每次做类似的故障诊断项目,都把导出和推理验证放在训练流程里当成正式步骤走一遍,不会再只盯着训练脚本跑完就收工。故障诊断最终要面对的是真实设备上的未知信号,训练环境里再好看的准确率,都比不上在两条真实采集的数据上跑一遍推理来得踏实。希望这份源码和这篇拆解,能帮你在毕业设计里少走一段路。
本文还有配套的精品资源,点击获取