简介基于RNN模型的轴承故障检测项目是一套完整的深度学习源码源自比赛项目面向机械故障诊断任务利用循环神经网络对轴承振动信号进行建模从而判断设备是否异常。该项目适合人工智能、通信工程、自动化、软件工程等专业的学生、老师或企业员工既可作为课程设计与毕业设计的选题也可用于入门深度学习的练手项目。压缩包共包含8个文件其中有3个Python源码文件、2个数据集压缩包、1个CSV数据表以及docx和md格式的说明文档整体大小约47.1MB。源码模块划分清晰get_data.py实现数据读取、清洗、归一化并划分训练集与测试集model.py定义RNN模型结构main.py负责训练和检测的完整流程。读者可以按序运行快速复现轴承故障检测实验并以此为基础调整网络参数或替换其他模型。该资源已有851人学习下载代码经过实际运行测试曾在答辩评审中获得94.5分的平均分质量可靠适合作为故障诊断项目的参考起点。1. 轴承故障检测为什么要用 RNN从振动信号里读时间模式在某次工业数据比赛里A 同学用 CNN 把轴承故障分类做到 91% 后死活上不去我换了一条 RNN 路线直接把结果推到 96%。听完你可能会愣一下——轴承故障检测不是图像任务吗为什么不用卷积问题恰恰出在这轴承振动信号是时间序列不是图像。外圈故障每过一个缺陷点就撞一次内圈故障的冲击随载荷位置周期性变化滚动体故障更是在保持架转动频率上做幅度调制。这些周期性冲击模式本质上是“前面几毫秒的信号能预测后面几毫秒的信号”而 RNN 按时间步展开做前向传播结构上就是干这个的。本文要拆解的这份“基于 RNN 模型的轴承故障检测 python 源码数据集(比赛项目).zip”核心思路就是把测点采集到的振动加速度信号切成窗口用循环结构逐个时间步读入最后用 softmax 分类出正常、外圈、内圈、滚动体故障等状态。适合正在做设备故障诊断的算法工程师、参加工业 AI 比赛的学生以及准备上预测性维护但对模型选型还没底的技术负责人。读完你不仅能跑通这套代码还能知道每个参数为什么设成这个值以及最坑的几步操作长什么样。2. RNN 模型的原理振动信号为什么适合用循环结构吃进去2.1 轴承振动信号的结构特征与 RNN 的契合点轴承振动信号和自然语言序列在统计特性上完全不同但有一个关键共性相邻时间点之间存在强依赖关系。外圈故障的特征频率大概在几十赫兹到几百赫兹之间而加速度传感器的采样率动辄 12kHz 甚至 20kHz。这意味着一个故障冲击事件横跨成百上千个采样点冲击之后还跟着一串衰减振荡。如果你把整个序列抻平丢给全连接网络模型很难学到“第 300 个采样点附近的振荡是第 150 个采样点处冲击的余波”这种关系。CNN 用卷积核在时间轴上滑动能捕捉局部模式但感受野的长度受卷积核尺寸和层数限制对于跨几百个点的故障周期需要堆很多层才能把依赖范围拉长。RNN 则不同它在每个时间步维护一个隐状态 h_t更新公式是h_t tanh(W_xh · x_t W_hh · h_{t-1} b_h)也就是说当前时刻的输出不仅取决于当前输入 x_t还取决于上一个时刻的隐状态 h_{t-1}。隐状态就是一个不断滚动的信息摘要它把“截止到当前时刻的振动模式”编码成一个固定维度的向量。往前的依赖能传多远取决于训练是否解决了梯度消失问题。对轴承故障信号来说故障冲击的周期在几十到几百个采样点之间这个依赖长度正好是 RNN 类模型擅长的区间。还有一个现实优势RNN 的参数量远小于同规模的全连接网络因为不同时间步共享同一套权重。在比赛场景下训练集通常只有几万到几十万个窗口参数太多反而容易过拟合。2.2 比赛源码里用的其实是 LSTM 或 GRU别被标题里的“RNN”字眼带偏。纯 vanilla RNN 在序列长度超过几十步之后梯度回传会遇到连乘效应要么梯度爆炸、要么梯度消失。轴承振动信号一个窗口就是 1024 或 2048 个点如果直接上 vanilla RNN训练到第 20 个 epoch 你就会发现 loss 开始震荡准确率卡住不动这是典型的梯度消失症状。所以比赛项目里真正落地的主力是 LSTM 或 GRU。LSTM 引入了输入门、遗忘门、输出门三个门控用细胞状态 c_t 作为信息传递的主干道遗忘门决定上一时刻的哪些信息要丢弃输入门决定新信息写入多少。GRU 更精简只有更新门和重置门两个门参数量比 LSTM 少大约四分之一在训练数据量有限时不容易过拟合训练速度也更快。对比项 | LSTM | GRU 门控数量 | 3 个门 | 2 个门 参数量 | 大 | 比 LSTM 少约 25% 训练速度 | 较慢 | 较快 小数据集表现 | 容易过拟合 | 更稳健 工业故障诊断场景 | 数据充足时可用 | 更常用我做这类项目默认先跑 GRU。比赛有 deadline你不可能花两周时间去调一个三层 LSTMGRU 单层起步效果不够再往上加。源码里如果看到nn.GRU或nn.LSTM的调用不要觉得它名不副实——RNN 是家族名LSTM/GRU 是具体实现。2.3 输出层与损失函数softmax 加交叉熵别用均方误差轴承故障检测的标签是互斥的一个窗口要么是正常要么是外圈故障不可能同时是两种。输出层用 softmax把模型的原始输出 logits 转成每个类别的概率分布概率之和为 1。损失函数用交叉熵cross-entropy loss它衡量的是预测分布和真实标签分布之间的差距梯度在错误分类时较大模型收敛快。初学者最容易犯的错是拿均方误差MSE当损失函数。分类任务的标签做完 one-hot 后是 0 和 1 组成的向量如果用 MSE模型会倾向于输出一个“中庸”的概率分布而不是 sharp 的预测——比如让每个类别的概率都接近 0.25这样 MSE 反而更小但分类任务的需求是要模型“有把握”地押注某一个类别。交叉熵不存在这个问题它只在预测概率和真实标签匹配时最小化。所以当你看到训练代码里用的是nn.CrossEntropyLoss()这个选择是对的别手痒改成nn.MSELoss()。3. python 源码与数据集落地把原始振动信号变成监督学习样本3.1 数据集构成与标签体系先从目录结构把信号和标签对上一份比赛项目的压缩包解压之后通常有源码目录和数据集目录。数据集目录的常见组织方式是按故障类型分文件夹每个文件夹里放对应工况下的原始信号文件也可能所有信号在一个大文件夹里文件名带标签后缀。第一步永远是遍历目录把每个文件的完整路径和它对应的标签建立映射关系。import os from collections import Counter LABEL_MAP { 正常: 0, 外圈: 1, 内圈: 2, 滚动体: 3, } def build_label_index(data_dir): index [] for root, dirs, files in os.walk(data_dir): for f in files: if f.endswith(.csv) or f.endswith(.npy): full_path os.path.join(root, f) # 从路径中的目录名或文件名前缀解析标签 folder_name os.path.basename(root) if folder_name in LABEL_MAP: label LABEL_MAP[folder_name] index.append((full_path, label)) # 打印类别分布检查有没有类别数据量特别少 labels [label for _, label in index] print(Counter(labels)) return index # 用法示例把数据集目录传进来 index build_label_index(./bearing_dataset)这段代码做的事情很笨但很有必要先扫描全目录再用文件夹名做标签解析最后统计类别分布。Counter(labels)输出形如{0: 2000, 1: 1800, 2: 1900, 3: 1200}的结果如果某个类别的样本数和其他类别差了一个数量级后续要做类别平衡处理否则模型会偏向样本量大的类。注意这里没有直接在遍历时加载信号数据因为比赛数据集动辄几个 GB一次性全部读进内存可能直接把机器弄死先建立索引后续按需加载更稳妥。3.2 滑动窗口切分与训练/验证集划分最容易踩的泄漏坑就在这里原始振动信号是连续采集的一维数组一个文件可能包含几百万个采样点对应几十秒甚至几分钟的连续运行数据。RNN 需要的输入是一个窗口也就是从长序列里截取的一段固定长度的片段。切窗口用滑动窗口法两个关键参数是窗口长度window_size和步长step。窗口长度要覆盖至少一个旋转周期以电机转速 1500rpm 为例一个周期等于 0.04 秒如果采样率是 12kHz一个周期约 480 个采样点取 1024 或 2048 都是安全的选择。太小的话窗口里可能捕捉不到完整的冲击模式太大则增加计算量且没必要。步长一般取窗口的一半做重叠这样相邻窗口有 50% 的重复区域能有效增加训练样本量同时保证故障冲击不会被切在窗口边缘而漏掉核心信息。import numpy as np def sliding_window_slice(signal, window_size1024, step512): num_windows max(0, (len(signal) - window_size) // step 1) windows [] for i in range(num_windows): start i * step end start window_size windows.append(signal[start:end]) return np.array(windows) # 示例把某个文件的信号切成窗口 raw_signal np.loadtxt(./bearing_dataset/外圈/sample_01.csv) windows sliding_window_slice(raw_signal, window_size1024, step512) print(windows.shape) # 比如输出 (234, 1024)这里有个决定成败的问题划分训练集和验证集时是按“窗口”随机划分还是按“文件”划分答案是按文件划分。原因是同一段连续信号切出的相邻窗口之间高度相关前一个窗口的最后 512 个点正好是后一个窗口的前 512 个点如果这些窗口同时出现在训练集和验证集里验证集的性能会被这种重叠信息推高造成模型“看过答案”的假象。比赛中这叫数据泄漏是排名虚高、线下验证失效的头号原因。我通常的做法是先把文件列表按标签分层划分成训练文件集合和验证文件集合每个集合内部再各自做滑窗。3.3 标准化与标签编码统计量只能从训练集算振动信号的幅值范围因测点位置和传感器灵敏度而异有的文件幅值在正负 5 之间有的在正负 50 之间。如果不做标准化模型对幅值尺度敏感GRU 单元里的激活函数容易饱和或进入线性区训练不稳定。标准化的标准做法是用训练集的均值和标准差把数据变换为均值为 0、标准差为 1 的分布。from sklearn.preprocessing import StandardScaler # 假设 train_windows 是训练集所有窗口拼成的大矩阵 (n_samples, window_size) n_samples, window_size train_windows.shape # 注意StandardScaler 是按列每个特征做标准化 # 这里我们只有一个测点但序列中不同位置的幅值分布都视为同分布 scaler StandardScaler() train_windows_flat train_windows.reshape(-1, 1) scaler.fit(train_windows_flat) # 只 fit 训练集的统计量 train_windows_norm scaler.transform(train_windows_flat).reshape(n_samples, window_size) # 验证集用同一个 scaler 做 transform不重新 fit val_windows_flat val_windows.reshape(-1, 1) val_windows_norm scaler.transform(val_windows_flat).reshape(val_windows.shape[0], window_size)很多人习惯在切窗之前对整个数据集做标准化也就是把训练集和验证集一起算均值和方差这样做的问题在于验证集的统计信息在训练阶段就已经被模型间接知道了相当于验证集的“考试题目”在考前泄了一半。严格的数据处理流程是先切窗再对训练集fit然后用训练集的统计量对验证集transform。标准化和切窗的顺序也值得注意——如果先标准化再切窗窗口之间的偏差可能有偏差但影响较小先切窗再标准化是更干净的做法。标签编码则简单得多把字符串标签映射成整数即可例如正常0, 外圈1, 内圈2, 滚动体3。LabelEncoder或手动字典都可以但手动字典更可控尤其是当数据集中还有“未知故障”这类类别时显式映射能避免排序导致的标签错位。4. 训练闭环从数据加载到故障分类的核心代码与参数设定4.1 模型定义单层 GRU 加分类头的简洁结构比赛源码里的模型结构通常不复杂核心是一个单层 GRU 加一个全连接分类头。输入张量的形状是(batch_size, seq_len, input_dim)其中seq_len是窗口长度input_dim是每个时间步的特征维度。如果只用单一测点的原始振动信号input_dim1如果把驱动端和风扇端的信号拼接起来input_dim2。import torch import torch.nn as nn class BearingGRUClassifier(nn.Module): def __init__(self, input_dim1, hidden_size32, num_layers1, num_classes4, dropout0.2): super().__init__() self.gru nn.GRU( input_sizeinput_dim, hidden_sizehidden_size, num_layersnum_layers, batch_firstTrue, dropoutdropout if num_layers 1 else 0.0, ) self.dropout nn.Dropout(pdropout) self.fc nn.Linear(hidden_size, num_classes) def forward(self, x): # x: (batch_size, seq_len, input_dim) _, h_n self.gru(x) # h_n: (num_layers, batch_size, hidden_size) # 取最后一层的隐状态作为序列表示 out h_n[-1] # (batch_size, hidden_size) out self.dropout(out) logits self.fc(out) # (batch_size, num_classes) return logits model BearingGRUClassifier(input_dim1, hidden_size32, num_layers1, num_classes4) print(model)hidden_size是 GRU 隐状态的维度决定了记忆容量。轴承故障模式相对清晰32 个节点足够表达大多数工况下的特征设到 128 也不会明显提升准确率反而增加过拟合风险和训练时间。num_layers1是因为单层 GRU 已经能捕捉几十到几百个时间步的依赖关系多层堆叠带来的收益在比赛这种中等问题规模上不显著。dropout0.2加在全连接层之前是对分类头做正则化。之所以不给单层 GRU 内部加 dropout是因为 PyTorch 要求num_layers 1时 GRU 层间 dropout 才生效。4.2 训练循环与优化器配置adam 加学习率衰减加早停训练配置有一个比较稳的组合优化器用 adam初始学习率1e-3batch size 取 64最大训练轮数 50配合早停early stopping机制当验证集准确率连续多个 epoch 不再刷新最优值时停止训练并回滚到最优权重。为什么用 adam振动信号经标准化后分布相对规整adam 的自适应学习率机制对这种场景很友好不需要手动精细调整学习率。SGD 虽然理论上能收敛到更好的局部最优点但对学习率和动量参数敏感比赛周期内不值得跟它较劲。import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset train_dataset TensorDataset( torch.tensor(train_windows_norm, dtypetorch.float32).unsqueeze(-1), torch.tensor(train_labels, dtypetorch.long), ) val_dataset TensorDataset( torch.tensor(val_windows_norm, dtypetorch.float32).unsqueeze(-1), torch.tensor(val_labels, dtypetorch.long), ) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, drop_lastTrue) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3) scheduler optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemax, factor0.5, patience3) best_val_acc 0.0 patience 0 best_model_state None for epoch in range(50): model.train() total_loss 0.0 for batch_x, batch_y in train_loader: optimizer.zero_grad() logits model(batch_x) loss criterion(logits, batch_y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() total_loss loss.item() model.eval() correct 0 total 0 with torch.no_grad(): for batch_x, batch_y in val_loader: logits model(batch_x) preds torch.argmax(logits, dim1) correct (preds batch_y).sum().item() total batch_y.size(0) val_acc correct / total scheduler.step(val_acc) print(fEpoch {epoch1}, Loss: {total_loss / len(train_loader):.4f}, Val Acc: {val_acc:.4f}) if val_acc best_val_acc: best_val_acc val_acc best_model_state {k: v.clone() for k, v in model.state_dict().items()} patience 0 else: patience 1 if patience 5: print(fEarly stopping at epoch {epoch1}, best val acc: {best_val_acc:.4f}) break model.load_state_dict(best_model_state)drop_lastTrue是为了避免最后一个 batch 样本数过少导致梯度方向不稳定。clip_grad_norm_做梯度裁剪防止 GRU 在回传过程中出现梯度爆炸——这是 RNN 类模型最经典的训练不稳定的来源。ReduceLROnPlateau在验证准确率连续 3 个 epoch 不提升时把学习率减半这是一个温和的调整策略比固定学习率硬跑 50 轮效果更稳。早停的 patience 设为 5也就是容忍验证集连续 5 个 epoch 不创新高。比赛数据集通常规模不大训练一个 epoch 只要一两分钟这套配置可以在 20 分钟内跑完一轮完整实验。4.3 评估与混淆矩阵别只看准确率要看到谁在拖后腿比赛排名看的是总体分类准确率但调试阶段只看准确率会漏掉关键信息。比如外圈故障和内圈故障的特征频率相近模型很可能把两者混在一起滚动体故障样本量少模型可能直接放弃学习它。这时需要看混淆矩阵。from sklearn.metrics import confusion_matrix, classification_report import seaborn as sns import matplotlib.pyplot as plt model.eval() all_preds [] all_labels [] with torch.no_grad(): for batch_x, batch_y in val_loader: logits model(batch_x) preds torch.argmax(logits, dim1) all_preds.extend(preds.numpy()) all_labels.extend(batch_y.numpy()) cm confusion_matrix(all_labels, all_preds) print(classification_report(all_labels, all_preds, target_names[正常, 外圈, 内圈, 滚动体])) plt.figure(figsize(6, 5)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabels[正常, 外圈, 内圈, 滚动体], yticklabels[正常, 外圈, 内圈, 滚动体]) plt.xlabel(Predicted) plt.ylabel(True) plt.show()classification_report会输出每个类别的精确率、召回率、F1-score。如果“滚动体”的召回率明显低于其他类别说明模型没有学会滚动体故障的判别特征需要检查这个类别的训练样本量是否足够或者考虑对少数类做加权采样。这是模型调试的重要步骤。5. 避坑清单RNN 轴承故障检测常见的四个翻车现场5.1 验证集性能虚高同源窗口泄漏现象训练过程中验证集准确率一路上涨到 98%你满心欢喜去提交测试集结果只有 86%线下线上差距大得离谱。原因验证集里的窗口和训练集里的窗口来自同一段连续信号相邻窗口重叠严重验证集已经“见过”训练数据的内容。解决按文件而不是按窗口划分训练/验证集。具体做法是先把文件列表按标签分层随机分成两份再对两个集合分别切窗。每类文件数量不要悬殊比如外圈文件有 20 个滚动体只有 6 个按文件划分时要用分层抽样保证每类在训练集和验证集中的比例一致。5.2 loss 震荡不收敛学习率过高或数据没标准化现象训练 loss 曲线像锯齿时而下降时而又跳回高位50 个 epoch 训完准确率还在 70% 附近晃悠。原因学习率太大GRU 的梯度方向在时序依赖下波动剧烈参数在最优解附近来回穿越另一种可能是输入信号没做标准化幅值分布过宽激活函数的输入落在饱和区。解决先把学习率降到1e-3以下如果还震荡就继续降到3e-4。同时检查数据预处理管线确认标准化用的是训练集的均值和标准差。5.3 过拟合但早停没触发验证集合太小导致噪声主导现象训练 loss 持续下降逼近 0验证准确率在第 10 个 epoch 后就停在一个值附近上下抖动早停机制因为抖动不断刷新“最优”而一直没触发白白跑了 50 轮。原因验证集的窗口数量太少比如每类只有 10 个窗口验证准确率的波动主要由随机误差主导无法真实反映模型泛化能力。解决增加验证集规模保证每类至少有 50 到 100 个窗口如果数据总量有限可以考虑以文件为单位做 K 折交叉验证而不是简单的单次划分。此外不要只依赖早停要同时盯训练 loss 和验证 loss 的差距差距持续拉大就是过拟合信号。5.4 标签与信号错位排序规则不一致导致的错乱现象某个类别的精确率异常高另一个类别又异常低而且怎么调参都恢复不了正常水平。原因加载数据时对文件名做了排序但不同操作系统下文件名的排序规则可能不一致信号文件和标签文件的顺序没对齐导致训练时模型把 A 类信号当成了 B 类标签。解决在数据加载完成后做一次可视化抽查。import matplotlib.pyplot as plt def plot_sample_signal(windows, label, sample_idx): plt.figure(figsize(10, 3)) plt.plot(windows[sample_idx]) plt.title(fLabel: {label}) plt.show() # 从验证集中抽一个样本看信号波形是否符合该故障类型预期 sample_windows val_windows_norm sample_labels val_labels plot_sample_signal(sample_windows, sample_labels[0], sample_idx0)外圈故障的时域波形通常在周期性位置出现明显的冲击尖峰正常信号的波形相对平稳。画出来看一眼就能发现标签是不是对得上。我每次加载新数据集都会抽三五个样本做可视化这个习惯帮我至少避开两次标签错位的坑。6. 进阶调优从比赛名次到真实产线部署的最后一公里比赛模型跑通了不等于产线能用。真实工业场景里测点数量、数据分布和故障定义都更复杂。如果你想让这套 RNN 方案从“比赛项目”升级成“可部署系统”可以从三个方向入手。第一是多传感器融合驱动端和风扇端信号拼成input_dim2的输入GRU 在每个时间步同时读两个测点的值这让模型能利用不同测点对同一故障的响应差异。具体做法就是修改模型输入的最后一个维度把两个测点的窗口在通道维度上拼接。第二是注意力机制GRU 输出每个时间步的隐状态而不是只取最后一步然后让注意力模块学习每个时间步的重要性权重再对隐状态做加权平均。这个做法能让模型意识到故障冲击出现的瞬间比其他平稳段更有判别价值通常能再提升 1 到 2 个百分点。第三是置信度阈值与误报漏报权衡生产环境里漏报一个轴承故障可能意味着整条产线停机代价远高于误报一个正常样本。我一般会在 softmax 输出的概率中加入一个“低置信拒绝”规则——如果最高类别的概率低于 0.8就输出“未知”并要求人工复核而不是硬性归到某个故障类别。这个阈值用验证集调画出 precision-recall 曲线后取一个误报和漏报平衡的点。推理时延方面GRU 的参数量小单层 32 隐节点只有几千个参数量化到 float16 之后可以直接跑在嵌入式设备上把窗口长度从 2048 缩到 512 还能再降低延迟不过要重新验证准确率。我现在的习惯是拿到任何竞赛项目源码先花半天把数据可视化做扎实再开始碰模型。数据没问题模型的结果才有解释的底气。希望帮到你。本文还有配套的精品资源点击获取