基于机器学习的机械故障诊断实战:从train.py到模型评估的完整链路
简介这份资源面向工业设备健康管理与故障诊断方向的学习者与工程人员聚焦如何用机器学习方法从振动、声音等传感器信号中识别异常模式并提前预警故障。压缩包共38个文件以20个py源码与18个pyc编译文件为主整体约75KB涵盖数据预处理、特征工程、模型训练与评估等环节的代码实现便于直接运行与二次修改。内容围绕频谱特性、时域统计量等特征提取以及决策树、随机森林、支持向量机、CNN、RNN等模型的选型与训练展开并涉及训练集、验证集、测试集划分与交叉验证、准确率、召回率、F1分数、AUC-ROC等评估指标同时讨论误报与漏报代价权衡及模型部署与实时监控思路。已有618人学习下载适合希望把机器学习落地到机械故障诊断场景的读者参考实践。1. 拆开这个故障诊断包为什么我建议先跑通 train.py 再谈算法选型工业设备维护这行有个反直觉的现象很多团队花大价钱上了振动传感器和采集卡数据攒了好几年真正能跑出故障预警的模型却没几个。问题往往不在算法本身而在于从原始信号到可训练特征之间那条链路没人认真搭过。这个「基于机器学习的机械故障诊断.zip」就是冲着这条链路来的——它把Machine_Learning_For_Fault_Diagnosis-main整个工程目录、dataset示例数据、model模型定义和train.py训练入口打包在一起让你能在一台普通笔记本上把「数据加载→特征处理→模型训练→指标评估」完整走一遍。适合刚接触机器学习故障诊断的工程师、做设备健康管理方向的学生以及想验证自己特征工程思路是否靠谱的一线维护人员。它不承诺开箱即用的工业级精度但能帮你把流程跑通、把坑踩明白。2. 从 dataset 到 train.py工程目录拆解与数据流梳理2.1 目录结构里藏着作者的设计意图拿到压缩包解压后根目录是Machine_Learning_For_Fault_Diagnosis-main下面几个关键路径值得先看清楚。dataset文件夹通常放的是振动信号或电流信号的原始数据格式可能是.csv、.mat或.txt具体取决于作者采集时用的设备。model目录里一般会有模型定义文件比如cnn_model.py、svm_model.py或random_forest.py每个文件对应一种算法实现。train.py是训练入口负责串联数据加载、预处理、模型构建和训练循环。有些版本还会带utils.py或preprocess.py把特征提取和归一化逻辑单独抽出来。我一般拿到这类工程包第一件事不是急着python train.py而是先看train.py的 import 部分。如果它从dataset直接读文件路径说明数据加载逻辑是硬编码的换数据集时需要改路径如果它通过参数或配置文件读取那移植性就好很多。这个包从命名看属于前者所以你得先确认dataset里的文件命名规则和train.py里写的路径是否一致。2.2 数据加载与预处理的关键参数假设dataset里是常见的 CWRU 轴承振动数据格式每个文件对应一种故障类型和负载条件。train.py里通常会有类似下面的加载逻辑import os import numpy as np import pandas as pd from scipy.io import loadmat # 数据目录根据实际解压路径调整 DATA_DIR ./dataset # 故障类别标签映射不同数据集需要改这里 LABEL_MAP { normal: 0, inner_race: 1, outer_race: 2, ball: 3 } def load_signal_files(data_dir): signals [] labels [] for fname in os.listdir(data_dir): if not fname.endswith(.mat): continue # 从文件名解析故障类型这里假设文件名包含类别关键词 for key, label in LABEL_MAP.items(): if key in fname.lower(): mat_data loadmat(os.path.join(data_dir, fname)) # DE 表示驱动端加速度信号不同数据集通道名不同 signal mat_data[DE].flatten() signals.append(signal) labels.append(label) break return np.array(signals), np.array(labels)这段代码的逻辑很直白遍历dataset下所有.mat文件根据文件名里的关键词判断故障类别然后读取DE通道的振动信号。参数上你要注意三个地方。第一LABEL_MAP的键必须和实际文件名匹配如果文件名是InnerRace_0hp.mat而你的映射写的是inner_race那就匹配不上样本会被静默跳过。第二mat_data[DE]这个键名取决于数据集版本CWRU 数据里驱动端信号通常叫DE风扇端叫FE基座叫BA用错通道等于换了个物理量。第三flatten()把二维数组压成一维如果原始信号是多通道的这里会丢失通道间关系做 CNN 时可能需要保留二维结构。2.3 特征工程时域统计量与频域特征的取舍原始振动信号直接喂给模型不是不行但大多数传统机器学习算法在手工特征上表现更稳。这个包里如果带了preprocess.py大概率会计算时域统计量均值、方差、峭度、峰值因子和频域特征FFT 幅值谱、功率谱密度。我一般会先跑一遍特征提取看看特征矩阵的维度是否合理。from scipy.stats import kurtosis, skew from scipy.fft import fft def time_domain_features(signal): features { mean: np.mean(signal), std: np.std(signal), kurtosis: kurtosis(signal), skewness: skew(signal), peak: np.max(np.abs(signal)), rms: np.sqrt(np.mean(signal ** 2)) } return features def freq_domain_features(signal, fs12000): n len(signal) yf fft(signal) xf np.linspace(0, fs / 2, n // 2) amplitude 2.0 / n * np.abs(yf[:n // 2]) # 取前 100 个频率点的幅值作为特征 return amplitude[:100]时域特征计算快、物理意义明确峭度对冲击性故障特别敏感RMS 反映能量水平。频域特征能捕捉周期性成分但维度高直接拼接会导致特征矩阵稀疏。常见做法是先用 FFT 找到主频带再计算该频带的能量占比而不是把整个幅值谱都塞进去。这个包如果用的是全谱拼接训练时容易过拟合你可以在train.py里把特征维度打印出来看看超过 200 维就要警惕了。3. 模型选型与训练SVM、随机森林和 CNN 在这个包里怎么用3.1 传统模型与深度学习模型的适用边界这个包的model目录如果同时提供了 SVM、随机森林和 CNN 的实现那作者的意图很明显让你对比不同算法在同一份数据上的表现。SVM 在小样本、特征维度适中时表现稳定尤其是线性核和 RBF 核调参主要看C和gamma。随机森林对特征尺度不敏感能输出特征重要性适合做初步的特征筛选。CNN 需要足够的样本量和合理的输入维度如果dataset里只有几百个样本CNN 很容易过拟合这时候不如用 SVM 或随机森林。我一般会按这个顺序试先跑随机森林看特征重要性排序把贡献低的特征去掉再用 SVM 在精简特征上训练对比准确率变化最后如果样本量超过两千再上 CNN 看能不能自动学到更复杂的模式。这个包如果train.py里默认用的是 CNN而你的数据只有几百条那第一轮跑出来准确率大概率在 60% 上下晃别急着怀疑代码先换模型。3.2 训练脚本的关键参数与修改点train.py里通常会有数据集划分、模型初始化、训练循环和评估输出。下面是一个典型的结构from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler from sklearn.svm import SVC from sklearn.metrics import classification_report, confusion_matrix # 加载数据和特征 signals, labels load_signal_files(DATA_DIR) features np.array([list(time_domain_features(s).values()) for s in signals]) # 标准化SVM 对尺度敏感必须做 scaler StandardScaler() features_scaled scaler.fit_transform(features) # 划分训练集和测试集stratify 保证类别比例一致 X_train, X_test, y_train, y_test train_test_split( features_scaled, labels, test_size0.3, random_state42, stratifylabels ) # SVM 训练C 和 gamma 需要根据数据调整 clf SVC(kernelrbf, C10, gammascale, probabilityTrue) clf.fit(X_train, y_train) # 评估 y_pred clf.predict(X_test) print(confusion_matrix(y_test, y_pred)) print(classification_report(y_test, y_pred))参数上你要盯住几个地方。test_size0.3是常见划分比例但如果某类故障样本特别少比如只有 20 条按 30% 划分后测试集只剩 6 条评估结果波动会很大这时候应该用交叉验证而不是单次划分。stratifylabels保证训练集和测试集的类别比例一致不加这个参数极端情况下测试集可能缺某一类。C10是 SVM 的惩罚系数值越大对误分类容忍度越低容易过拟合值太小则欠拟合。gammascale是自动根据特征方差设的如果特征维度差异大可以手动设成1 / (n_features * X.var())。3.3 训练过程中的监控与日志train.py如果只是跑完打印一个准确率那调试起来会很痛苦。我一般会加几行日志把每轮的损失、验证集准确率和学习率打出来。如果是 CNN用 PyTorch 的话可以这样写import torch import torch.nn as nn # 假设 model 已经定义好criterion 和 optimizer 也初始化了 for epoch in range(num_epochs): model.train() running_loss 0.0 for inputs, targets in train_loader: optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, targets) loss.backward() optimizer.step() running_loss loss.item() # 每 5 轮打印一次训练损失和验证准确率 if (epoch 1) % 5 0: model.eval() correct 0 total 0 with torch.no_grad(): for inputs, targets in val_loader: outputs model(inputs) _, predicted torch.max(outputs, 1) total targets.size(0) correct (predicted targets).sum().item() val_acc correct / total print(fEpoch {epoch1}, Loss: {running_loss/len(train_loader):.4f}, Val Acc: {val_acc:.4f})这段代码的关键在于把训练损失和验证准确率分开看。如果训练损失持续下降但验证准确率不升反降那就是过拟合了需要加 Dropout、减小模型规模或增加数据增强。如果训练损失震荡厉害可能是学习率太大试着除以 10。这个包如果没带验证集划分逻辑你得自己从训练集里再切一部分出来别直接用测试集调参那样评估结果会偏乐观。4. 避坑与排查数据泄漏、标签错位和过拟合的典型翻车现场4.1 现象测试集准确率 99%换一批数据掉到 50%原因数据泄漏。常见情况是特征提取时用了全局统计量比如对整个数据集算均值和方差再做标准化而不是只在训练集上拟合StandardScaler。另一个来源是同一段信号被切分成多个样本后随机划分导致训练集和测试集里有高度重叠的片段。解决标准化只在训练集上fit测试集用同样的transform。信号切分时按时间段划分而不是随机打乱。如果数据来自同一台设备的连续采集最好按采集批次划分训练和测试。4.2 现象混淆矩阵里某一类永远预测不对原因标签错位。文件名解析逻辑有漏洞比如inner_race和ball的文件名里都包含race或ball的变体匹配时先命中了错误的键。或者LABEL_MAP的键值对写反了正常样本被标成故障。解决在load_signal_files里加一行打印把每个文件名和匹配到的标签输出人工核对前 20 条。如果文件名规则复杂改用正则表达式精确匹配别用简单的in判断。4.3 现象训练损失降到 0.01验证损失却一直在涨原因过拟合。模型参数太多、训练样本太少、或者特征维度远大于样本数。CNN 在没有数据增强的情况下几千个参数就能把几百个样本背下来。解决先减模型复杂度把 CNN 的卷积层从 5 层降到 3 层全连接层神经元减半。加 L2 正则化PyTorch 里在 optimizer 里设weight_decay1e-4。如果样本实在少用传统机器学习模型替代深度学习SVM 在小样本上的泛化能力通常更好。4.4 现象训练时 GPU 利用率不到 30%速度很慢原因数据加载成了瓶颈。DataLoader的num_workers设成 0或者预处理逻辑写在__getitem__里每次重复计算。解决把num_workers设成 4 或 8具体看 CPU 核数。预处理能提前做的就提前做比如 FFT 和特征提取在数据集构建阶段一次性算完存成.npy文件训练时直接读。4.5 现象换一台机器跑准确率差了好几个点原因随机种子没固定。train_test_split的random_state、PyTorch 的torch.manual_seed、NumPy 的np.random.seed只要有一个没设结果就不可复现。解决在train.py开头统一设种子import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False set_seed(42)注意cudnn.deterministic True会稍微降低训练速度但能保证卷积操作的结果可复现。5. 进阶技巧用交叉验证和混淆矩阵把模型性能压榨到边界5.1 交叉验证替代单次划分单次train_test_split的评估结果受随机性影响很大尤其是小样本。用 5 折交叉验证能给出更稳定的性能估计。sklearn的cross_val_score可以直接用from sklearn.model_selection import cross_val_score, StratifiedKFold # 分层 K 折保证每折的类别比例一致 cv StratifiedKFold(n_splits5, shuffleTrue, random_state42) scores cross_val_score(clf, features_scaled, labels, cvcv, scoringf1_macro) print(fF1 macro: {scores.mean():.4f} /- {scores.std():.4f})scoringf1_macro比准确率更适合故障诊断因为各类样本不均衡时准确率会被多数类主导。f1_macro对每个类别算 F1 再取平均少数类的表现也能反映出来。如果某类 F1 明显低于其他类说明模型对该类故障不敏感需要针对性补充样本或调整特征。5.2 混淆矩阵的细读方法confusion_matrix输出的矩阵对角线是正确分类非对角线是误判。我一般会把它转成百分比形式看每个类别的误判流向import seaborn as sns import matplotlib.pyplot as plt from sklearn.metrics import confusion_matrix cm confusion_matrix(y_test, y_pred) # 按行归一化显示每个真实类别的预测分布 cm_normalized cm.astype(float) / cm.sum(axis1)[:, np.newaxis] plt.figure(figsize(8, 6)) sns.heatmap(cm_normalized, annotTrue, fmt.2f, cmapBlues) plt.xlabel(Predicted) plt.ylabel(True) plt.title(Normalized Confusion Matrix) plt.show()如果inner_race的样本有 30% 被误判成ball说明这两个故障在特征空间里重叠严重。常见原因是特征提取时只用了时域统计量没有捕捉到冲击频率的差异。可以加一段包络谱分析或者用scipy.signal.envelope提取包络后做 FFT把包络谱的峰值频率作为额外特征。5.3 模型集成与阈值调整如果单模型性能卡在 85% 上不去可以试试集成。随机森林本身是集成方法但你可以把 SVM、随机森林和一个小型 CNN 的预测概率平均一下from sklearn.ensemble import VotingClassifier # 假设 svm_clf、rf_clf、cnn_clf 已经定义好 ensemble VotingClassifier( estimators[(svm, svm_clf), (rf, rf_clf), (cnn, cnn_clf)], votingsoft # 用预测概率投票比硬投票更稳 ) ensemble.fit(X_train, y_train)votingsoft要求每个基模型都能输出概率SVM 需要设probabilityTrueCNN 需要在最后加 Softmax。集成后如果准确率提升不明显说明基模型之间的差异性不够换一个完全不同类型的模型加入比如 KNN 或梯度提升树。另一个技巧是调整分类阈值。默认阈值是 0.5即预测概率超过 0.5 就判为故障。但在故障诊断中漏报的代价通常高于误报所以可以把阈值降到 0.3让模型更倾向于报故障。代价是误报增多需要根据实际维护成本权衡。# 获取正类故障的预测概率 y_prob clf.predict_proba(X_test)[:, 1] # 降低阈值到 0.3 y_pred_adjusted (y_prob 0.3).astype(int) print(classification_report(y_test, y_pred_adjusted))从那以后我每次拿到新的故障诊断包都强制先跑一遍交叉验证和归一化混淆矩阵再决定要不要调模型结构。很多所谓的「模型不行」其实是数据划分和评估方式出了问题。希望这个包能帮你把故障诊断的完整链路跑通少走一些我当年走过的弯路。本文还有配套的精品资源点击获取