基于一维CNN的轴承故障诊断：从振动信号到混淆矩阵的完整实战

简介这份资源面向深度学习与工业故障诊断方向的研究者和工程师提供一套可直接运行的完整方案帮助读者从传感器数据中自动识别设备异常状态免去从零搭建系统的成本。包内共63个文件以mat数据文件、py源码、xml配置、csv数据集为主另含h5训练权重、png模型结构图与readme说明压缩包约57.74MB。内容覆盖一维卷积与LSTM两条诊断路线包含数据预处理、模型定义、训练脚本、验证测试与预测接口等模块并附有轴承故障数据集与训练日志便于对照复现。已有11421人学习下载代码注释较为完整读者可据此理解时序与频域特征提取思路调整网络结构与参数并迁移到其他类型的故障诊断任务中兼具学术研究与工程实践参考价值。1. 从一段振动信号说起这套故障诊断代码到底能干什么设备还没停报警先响了。产线上的振动传感器每秒钟吐出几千个采样点靠人盯着频谱图找异常基本等于大海捞针。这套基于深度学习的故障诊断模型代码和数据解决的就是这件事把原始振动信号丢进去模型直接告诉你轴承是内圈故障、外圈故障还是滚动体故障甚至能区分不同损伤直径。它适合两类人——一类是刚接触工业信号、想跑通一个完整诊断流程的算法新手另一类是手里有设备数据、想找个能改能扩的baseline的工程师。代码覆盖了从数据加载、特征提取、模型训练到评估可视化的全链路数据集用的是公开的轴承故障数据目录结构清晰改起来不费劲。下面我按实际拆项目的顺序把这份资源从跑通到改造讲一遍。2. 项目结构与技术选型为什么是CNN振动信号2.1 目录长什么样每个文件干什么拿到一份代码我第一件事不是急着跑而是先看目录。这套资源的组织方式很典型适合做二次开发的起点。根目录下通常有这么几块fault_diagnosis/ ├── data/ # 原始数据与预处理后的数据 │ ├── raw/ # 原始振动信号文件.mat 或 .csv │ └── processed/ # 切分、归一化后的训练/测试集 ├── models/ # 模型定义 │ ├── cnn.py # 一维卷积网络主体 │ └── utils.py # 训练辅助函数 ├── configs/ # 超参数配置 │ └── default.yaml ├── train.py # 训练入口 ├── evaluate.py # 评估与混淆矩阵输出 ├── requirements.txt └── README.mddata/raw放的是原始信号一般是.mat格式每个文件对应一种故障类型在不同负载下的采样。data/processed是脚本切分后的结果通常按滑动窗口切成固定长度的样本比如 1024 个点一个样本。models/cnn.py是核心定义了一维卷积、池化、全连接的结构。configs/default.yaml把学习率、batch size、窗口长度这些参数抽出来改实验不用动代码。train.py和evaluate.py是两个入口一个负责训一个负责看结果。注意不同来源的代码在目录命名上会有差异但核心逻辑跑不出这几块。如果拿到手的版本没有configs参数大概率硬编码在train.py里改的时候要小心。2.2 技术选型一维CNN为什么比手工特征更省事传统故障诊断的做法是先做傅里叶变换再提取时域、频域统计量最后喂给SVM或随机森林。这套流程的问题在于特征工程高度依赖经验换个工况可能就得重新调。深度学习方案直接把原始信号或简单预处理后的信号送进网络让卷积核自己去学特征。选一维CNN而不是RNN或Transformer原因很实际。振动信号是典型的时间序列但故障特征往往体现在局部冲击上一维卷积核正好擅长捕捉这种局部模式。相比LSTMCNN训练更快参数更少在小样本工业数据集上不容易过拟合。相比TransformerCNN对数据量的要求低得多几百个样本就能训出不错的效果。常见做法是堆叠三到五层卷积每层后面接BN和ReLU最后用全局平均池化或展平后接全连接分类。import torch import torch.nn as nn class FaultCNN(nn.Module): def __init__(self, num_classes10, input_len1024): super().__init__() # 第一层卷积提取局部冲击特征 self.conv1 nn.Conv1d(1, 16, kernel_size15, stride2, padding7) self.bn1 nn.BatchNorm1d(16) # 第二层感受野扩大捕捉更复杂的模式 self.conv2 nn.Conv1d(16, 32, kernel_size7, stride2, padding3) self.bn2 nn.BatchNorm1d(32) # 第三层进一步抽象 self.conv3 nn.Conv1d(32, 64, kernel_size5, stride2, padding2) self.bn3 nn.BatchNorm1d(64) self.relu nn.ReLU() self.pool nn.AdaptiveAvgPool1d(1) # 全局平均池化抗过拟合 self.fc nn.Linear(64, num_classes) def forward(self, x): # x shape: (batch, 1, input_len) x self.relu(self.bn1(self.conv1(x))) x self.relu(self.bn2(self.conv2(x))) x self.relu(self.bn3(self.conv3(x))) x self.pool(x).squeeze(-1) # (batch, 64) return self.fc(x)这段代码里kernel_size的选择有讲究。第一层用15是为了覆盖一个冲击周期内的关键点后面逐层减小配合stride降采样。AdaptiveAvgPool1d(1)把每个通道压成一个数替代了传统的展平加全连接参数量少对输入长度变化也更鲁棒。num_classes根据数据集里故障类别数改常见轴承数据集是10类正常3种故障位置×3种损伤直径。input_len是切窗长度1024是常用值对应采样频率12kHz下约85毫秒的信号。2.3 数据预处理切窗、归一化、标签映射原始信号不能直接喂网络。第一步是切窗把长信号切成固定长度的样本。滑动窗口的步长决定样本数量步长小样本多但重叠高步长大样本少但可能漏掉瞬态。常见做法是步长取窗口长度的一半。import numpy as np def sliding_window(signal, window_size1024, step512): 将长信号切分为固定长度的样本 samples [] for start in range(0, len(signal) - window_size 1, step): samples.append(signal[start:start window_size]) return np.array(samples) def normalize(samples): 按样本做零均值单位方差归一化 mean samples.mean(axis1, keepdimsTrue) std samples.std(axis1, keepdimsTrue) 1e-8 return (samples - mean) / std归一化这里有个坑是按整个数据集归一化还是按每个样本归一化。按样本归一化能消除不同负载下的幅值差异让模型更关注波形形状工业场景下更常用。标签映射就是把文件夹名或文件名映射成整数比如normal对应0inner_race_007对应1以此类推。映射表最好存成json评估的时候反查用。3. 训练与评估从配置文件到混淆矩阵3.1 配置文件怎么改参数含义逐个说把超参数抽到yaml里是为了做实验的时候不用翻代码。一个典型的配置文件长这样data: window_size: 1024 step: 512 batch_size: 64 num_workers: 4 model: num_classes: 10 input_len: 1024 train: epochs: 50 lr: 0.001 weight_decay: 1e-4 device: cuda output: save_dir: checkpoints/ log_interval: 10window_size和step决定样本数量和重叠度。batch_size在显存允许的情况下尽量大64是安全值。lr用Adam的话0.001是默认起点如果loss震荡就降到0.0005。weight_decay加一点L2正则防止过拟合。device改成cpu也能跑就是慢。save_dir是模型权重保存路径训练脚本一般会保存验证集上最好的那个epoch。提示如果数据集类别不平衡比如正常样本远多于故障样本可以在损失函数里加weight参数或者用重采样。代码里如果没写自己加几行就行。3.2 训练循环损失函数、优化器、早停训练脚本的核心逻辑不复杂但有几个细节决定成败。损失函数用交叉熵优化器用Adam学习率调度用余弦退火或StepLR。早停是防止过拟合的后悔药验证集loss连续几个epoch不降就停。import torch import torch.nn as nn from torch.optim import Adam from torch.optim.lr_scheduler import CosineAnnealingLR def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0 for x, y in loader: x, y x.to(device), y.to(device) optimizer.zero_grad() logits model(x) loss criterion(logits, y) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(loader) # 初始化 model FaultCNN(num_classes10).to(device) criterion nn.CrossEntropyLoss() optimizer Adam(model.parameters(), lr0.001, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max50) best_val_loss float(inf) patience 10 counter 0 for epoch in range(50): train_loss train_one_epoch(model, train_loader, optimizer, criterion, device) val_loss evaluate(model, val_loader, criterion, device) scheduler.step() if val_loss best_val_loss: best_val_loss val_loss torch.save(model.state_dict(), checkpoints/best.pth) counter 0 else: counter 1 if counter patience: print(fEarly stop at epoch {epoch}) breakCosineAnnealingLR让学习率按余弦曲线下降前期大后期小比固定学习率更容易收敛到好的局部最优。patience10是经验值数据集小的话可以设5。保存best.pth而不是最后一个epoch的权重是因为过拟合之后验证集性能会下降。评估函数里记得把模型切到eval()模式关掉dropout和BN的更新。3.3 评估指标准确率之外还要看什么准确率在类别不平衡的时候会骗人。比如正常样本占90%全猜正常也有90%准确率。所以评估必须看混淆矩阵和每个类别的召回率。from sklearn.metrics import confusion_matrix, classification_report import seaborn as sns import matplotlib.pyplot as plt def evaluate(model, loader, criterion, device): model.eval() all_preds, all_labels [], [] total_loss 0 with torch.no_grad(): for x, y in loader: x, y x.to(device), y.to(device) logits model(x) loss criterion(logits, y) total_loss loss.item() preds logits.argmax(dim1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(y.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) print(classification_report(all_labels, all_preds)) plt.figure(figsize(10, 8)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues) plt.xlabel(Predicted) plt.ylabel(True) plt.savefig(confusion_matrix.png) return total_loss / len(loader)混淆矩阵能看出模型把哪两类搞混了。比如内圈故障和滚动体故障容易混说明特征区分度不够可以考虑加宽第一层卷积核或者增加一层。classification_report输出每个类别的precision、recall、f1-score重点关注recall低的类别。如果某个故障类别recall特别低要么是样本太少要么是这种故障的冲击特征被噪声淹没了。4. 避坑与排查跑不通的时候先看这几条4.1 现象loss不降准确率卡在10%左右原因通常是标签没对齐。切窗之后样本顺序和标签顺序错位或者标签映射表写错了。解决方法是打印前几个batch的标签和对应的文件名人工核对。另一个可能是输入没归一化信号幅值差异大导致梯度爆炸。检查normalize函数有没有被调用。4.2 现象训练集准确率很高验证集很低典型的过拟合。先看数据集大小如果每类只有几十个样本模型参数量要降。把第三层卷积的通道数从64降到32或者加dropout。另外检查训练集和验证集是不是来自同一个负载工况如果验证集是不同负载那属于域偏移问题需要加域适应模块不是简单调参能解决的。4.3 现象GPU显存不够报CUDA out of memory减小batch_size从64降到32甚至16。如果还不行把input_len从1024降到512但注意感受野要相应调整。另外检查num_workers是不是设太大了Windows下设0或2比较稳。用torch.cuda.empty_cache()清理缓存但根本办法还是减小模型或输入。4.4 现象评估结果每次跑都不一样没有固定随机种子。在训练脚本开头加import torch import numpy as np import random 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 Falsecudnn.deterministicTrue会让训练慢一点但结果可复现。如果追求速度可以关掉但对比实验的时候必须开。4.5 现象混淆矩阵里某一类全是0检查数据加载部分是不是某个类别的文件没被读到。常见原因是文件名过滤规则写死了比如只匹配*.mat但那个类别是.csv。或者类别映射字典里漏了某一类导致标签越界。打印dataset.classes和label_map核对。5. 二次开发换数据集、改模型、加注意力5.1 换成自己的数据需要动哪几个地方换数据集的核心是改数据加载部分。假设你手里的是CSV格式每行一个采样点文件名包含标签信息。需要改三处一是data/raw的读取逻辑把.mat读取换成pd.read_csv二是切窗函数如果采样频率不同window_size要按比例调整比如从12kHz换到48kHz窗口长度可以保持时间跨度不变设为4096三是标签映射根据你的文件名规则重写。import pandas as pd import os def load_custom_data(root_dir): samples, labels [], [] label_map {normal: 0, inner: 1, outer: 2, ball: 3} for fname in os.listdir(root_dir): if not fname.endswith(.csv): continue # 从文件名提取标签比如 inner_007.csv label_key fname.split(_)[0] if label_key not in label_map: continue signal pd.read_csv(os.path.join(root_dir, fname)).values.flatten() windows sliding_window(signal, window_size4096, step2048) samples.append(windows) labels.extend([label_map[label_key]] * len(windows)) return np.vstack(samples), np.array(labels)改完之后先跑一个mini batch打印样本形状和标签分布确认没问题再全量训练。如果类别不平衡在CrossEntropyLoss里加weighttorch.tensor([1.0, 2.0, 2.0, 2.0])这种给少数类更高权重。5.2 模型改造加注意力机制值不值一维CNN加注意力常见的是在卷积后面接SE模块或CBAM。SE模块的思路是让网络自动学习每个通道的重要性对故障诊断这种通道间差异明显的任务有帮助。实现很简单class SEBlock(nn.Module): def __init__(self, channels, reduction16): super().__init__() self.fc1 nn.Linear(channels, channels // reduction) self.fc2 nn.Linear(channels // reduction, channels) self.sigmoid nn.Sigmoid() def forward(self, x): # x: (batch, channels, length) w x.mean(dim2) # 全局平均池化 w self.sigmoid(self.fc2(nn.ReLU()(self.fc1(w)))) return x * w.unsqueeze(2)把SEBlock插在每层卷积之后、ReLU之前。reduction16是压缩比通道数少的时候可以设8。加注意力的代价是参数量和计算量略增但在小数据集上通常能涨1到3个百分点的准确率。如果数据集本身很大注意力带来的收益可能不明显不如直接加深网络。5.3 验证改造效果控制变量与消融实验改完模型别急着高兴先做消融。固定随机种子同一份数据划分只改一个变量。比如对比加SE和不加SE跑三次取平均。如果提升在1%以内可能是随机波动不值得增加复杂度。另外注意验证集不能用来调超参数否则会过拟合验证集。正确做法是划出训练集、验证集、测试集验证集用于早停和选模型测试集只在最后跑一次。我自己的习惯是每次改完模型先跑一个5 epoch的快速实验看loss曲线有没有异常再决定要不要跑完整训练。这样能省下大量等GPU的时间。从那以后我每次换数据集或者改网络结构都强制走一遍「打印样本形状→检查标签分布→跑mini batch→看loss曲线」这个流程翻车的次数少了很多。希望帮到你。本文还有配套的精品资源点击获取