简介基于深度学习的轴承故障诊断平台面向机械故障诊断与深度学习方向的毕业设计、课程设计开发者针对传统诊断高度依赖专业知识、效率低下、难以适应复杂工况等痛点提供一套可运行、可扩展的完整方案。压缩包共63个文件、大小14.85MB其中含11个Python脚本覆盖数据预处理、特征提取、模型训练、模型评估、故障诊断等关键环节10个MAT格式的轴承振动数据集用于模型训练与验证34张可视化结果图包含训练曲线、混淆矩阵、ROC曲线、精确率召回率曲线等方便直观对比效果。平台集成了CNN、LSTM、GRU、MCNN-LSTM等多种网络并引入YOLO进行实时故障检测同时提供了界面化的交互工具和标准化处理功能可帮助使用者快速跑通从原始信号到故障分类的完整链路。目录结构清晰、模块划分明确适合论文复现、课程展示或毕设参考。目前已有59人学习下载。1. 轴承故障诊断为什么值得用深度学习去做传统轴承故障诊断靠老师傅拿听音棒贴在轴承座上听经验老到的人能从异响里分辨出外圈剥落还是内圈点蚀但这份手艺很难复制判断标准也说不清道不明。深度学习出现以后这个场景被彻底改写了——振动信号转成时频图喂给卷积网络模型自己学习故障特征准确率能做到 98% 以上而且整套流程是可以复现、可以参数化的。这份《基于深度学习的轴承故障诊断平台》资源覆盖的就是从原始振动信号到最终故障类型输出的完整链路包括信号预处理、时频图生成、模型训练和平台部署。适合正在做毕业设计、课程设计的同学也适合想快速搭一套诊断 demo 的工程师下载下来照着跑一遍就能把原理和代码对上了。2. 先把信号变成图像STFT 时频图生成与样本集制作2.1 为什么不能把原始振动信号直接丢给卷积网络很多初学者拿到轴承振动数据后第一反应是把一维信号直接拼成二维矩阵输入 CNN。这个做法不算错但效果通常不理想。原因在于卷积核的平移共享特性更适合捕捉图像中的局部模式而原始振动信号的故障特征分散在时间轴和频率轴两个维度上时域波形里特征被背景噪声淹没直接堆成一维卷积网络需要非常深的网络才能学到有区分度的特征小数据集上根本训不动。工程上的常见做法是先做时频变换把一维信号变成二维时频图。时频图横轴是时间纵轴是频率颜色深浅代表能量强度轴承故障的周期性冲击会表现为特定频带上的能量聚集这个模式在二维图像里非常明显可以直接套用成熟的图像分类网络。常见选择是短时傅里叶变换 STFT计算量可控物理含义直观配合 Hann 窗和 50% 重叠率能拿到不错的频率分辨率。比小波变换好调参比 Wigner-Ville 分布抗交叉项干扰作为入门方案性价比最高。2.2 Python 实现 STFT 时频图批量生成import numpy as np from scipy.signal import stft from PIL import Image import os fs 12800 # 采样率单位 Hz常见工业采集卡的默认配置 window_len 256 # 窗长决定频率分辨率 hop_len 128 # 步长与窗长配合实现 50% 重叠 total_nums 300 # 待处理的原始信号文件数量 output_dir stft_images os.makedirs(output_dir, exist_okTrue) for idx in range(total_nums): # 假设 read_vibration_signal 读入一段长度为 fs 秒的振动信号 signal read_vibration_signal(idx) f, t, Zxx stft(signal, fsfs, npersegwindow_len, noverlaphop_len, windowhann) # Zxx 是复数矩阵取幅值后转分贝 spec np.abs(Zxx) spec_db 20 * np.log10(spec / spec.max() 1e-12) # 分贝值归一化到 0-255保存为灰度图 img_arr (spec_db - spec_db.min()) / (spec_db.max() - spec_db.min()) img_arr (img_arr * 255).astype(np.uint8) Image.fromarray(img_arr).save( os.path.join(output_dir, fsample_{idx:04d}.png) )这段代码的核心逻辑是把每段原始振动信号做 STFT得到复数矩阵后取幅值并转换成分贝刻度。之所以用分贝而不是线性幅值是因为轴承故障冲击产生的能量和正常振动能量的差距往往有几个数量级线性刻度下微弱故障特征会被淹没分贝刻度能把这些细微差别拉开。窗长 256 对应的频率分辨率是 50 Hz对轴承故障特征频率来说这个数值是足够的。如果你的数据采集频率更高比如 48 kHz建议把窗长提到 512 或 1024否则每个频率仓覆盖范围太宽特征会糊在一起。步长 128 配合窗长 256 实现了 50% 重叠这是 STFT 的常见配置时间分辨率足够计算量也可接受。2.3 样本集制作与数据增强的取舍时频图生成只是第一步真正决定模型上限的是样本集的构造方式。一个常见误区是每段原始信号只生成一张图整个数据集稀疏得很模型很容易过拟合。我的做法是把长信号切成固定长度的小段每段时长 1 秒重叠 50%这样一段 10 秒的原始信号能切出 19 段数据集直接扩到将近 20 倍。数据增强方面要克制。对时频图做随机翻转、随机裁剪这类图像增强技巧在自然图像上很有效但放在时频图上需要谨慎。时间轴方向的小幅平移可以接受因为故障冲击本来就有随机性但频率轴的翻转会破坏频率和故障类型的对应关系——内圈故障和外圈故障的特征频率不同翻转之后模型学到的是错误映射。经验做法是只做两类增强给原始信号叠加高斯白噪声后再生成时频图模拟不同工况下的背景噪声变化以及时间轴上的小幅裁剪。这两类增强都在信号层面做生成图像后不再做任何图像级增强。另外每类故障的样本数要尽量均衡不然训练出来的模型会对样本多的类别过度偏好这个在后面的避坑章节会细说。3. 模型选型与训练流程从 ResNet18 微调到故障分类3.1 一维 CNN 还是二维 CNN先看你的数据形态选模型之前先想清楚一个问题你手上最终的数据形态是什么。如果你把信号转成了时频图那就用二维 CNN这是主流选择因为可以直接加载 ImageNet 预训练权重做迁移学习小数据集也能收敛得很稳。如果你坚持用原始一维信号那就得走 1D-CNN 路线网络结构得自己设计卷积核尺寸、空洞率、池化策略都要反复试调试成本高出一截。这份资源里默认的路线是二维 CNN 加 ResNet18 迁移学习。选 ResNet18 而不是 ResNet50 或更深的网络是因为轴承故障诊断的数据集规模通常不大深度网络的参数量在这种数据量下容易过拟合ResNet18 的残差结构已经足够提取时频图上的纹理特征。而且 ResNet18 在 CPU 上推理速度快后续做成实时诊断平台时不会成为性能瓶颈。如果你的数据量确实很大比如每类故障样本超过一万张可以考虑升级到 ResNet50。但训练轮数要相应增加学习率要调低否则深层网络在小数据集上很容易训出训练集准确率极高、验证集一塌糊涂的结果。3.2 PyTorch 迁移学习训练脚本及参数清单import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from torchvision import models, transforms from PIL import Image num_classes 4 # 正常、外圈故障、内圈故障、滚动体故障 batch_size 32 lr 1e-4 epochs 50 device torch.device(cuda if torch.cuda.is_available() else cpu) class FaultDataset(Dataset): def __init__(self, img_paths, labels, transformNone): self.img_paths img_paths self.labels labels self.transform transform def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img Image.open(self.img_paths[idx]).convert(L) # 灰度图转三通道适配 ImageNet 预训练权重 img img.convert(RGB) if self.transform: img self.transform(img) label self.labels[idx] return img, label train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 加载预训练 ResNet18替换最后全连接层 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, num_classes) model model.to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lrlr, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxepochs) train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workers4) for epoch in range(epochs): model.train() running_loss 0.0 correct 0 total 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() running_loss loss.item() scheduler.step() train_acc 100.0 * correct / total print(fEpoch [{epoch 1}/{epochs}] Loss: {running_loss / len(train_loader):.4f} fTrain Acc: {train_acc:.2f}%)代码逻辑分几块数据集类负责读图、转灰度再转三通道这是为了适配 ImageNet 预训练权重的输入要求模型部分替换了 ResNet18 的最后一层全连接输出类别数改成你的故障类别数训练循环是标准的 PyTorch 流程每轮训练完更新学习率。迁移学习的处理方式值得注意默认情况下所有层的权重都会参与训练这在数据量较大时没问题。如果数据量很少比如几百张图建议冻结前面的卷积层只训练最后一层全连接代码里加一行model.requires_grad_(False)再把model.fc.requires_grad_(True)训练速度更快过拟合风险也更低。训练参数上AdamW 配合 1e-4 初始学习率算是个稳健组合weight_decay 设 1e-4 控制权重范数。余弦退火学习率调度器比固定学习率好前期大步探索后期小步收敛不容易出现 loss 震荡。3.3 训练过程的监控指标与模型保存策略训练过程不能只看训练集准确率我最关心的是验证集准确率和 loss 曲线的间距。训练集 loss 持续下降但验证集 loss 不降反升这是典型的过拟合信号。常见的应对手段有三个加大数据增强的强度、增加 dropout、降低模型容量。对迁移学习场景来说降低模型容量换成 ResNet18 就已经是低配了优先考虑前两种。模型保存有个血泪经验不要每轮都保存磁盘很快就满了。我的习惯是保存验证集准确率最高的那一轮同时只保留最近五个检查点代码里用一个列表记录准确率每次有新高就覆盖旧的。另外保存时把模型结构定义和权重一起打包用torch.save(model.state_dict(), path)就够了加载时先定义好网络结构再load_state_dict。训练完成后再拿测试集测一次泛化指标注意测试集不能参与过任何训练和验证流程否则指标没有参考意义。4. 避坑与常见问题排查训练跑崩和假高准确率的五个根源4.1 模型准确率虚高归一化统计量泄漏现象训练集准确率 99%测试集准确率也 99%但部署到现场数据上准确率掉到 60% 以下。原因对整批数据统一计算均值和标准差做归一化或者统一做了最大值最小值缩放到 0-255这会让归一化统计量同时包含训练集和测试集的信息相当于测试集的信息提前泄露给了模型训练过程。测试时新数据按训练集的统计量归一化分布对不上准确率自然崩。解决归一化统计量只能在训练集上计算然后把训练集的均值和标准差保存到文件测试时加载同一个文件使用。代码上就是transform里 Normalize 的 mean 和 std 用训练集算出来的值测试集和推理阶段用同样的值.4.2 时频图参数不一致现象训练时准确率很高保存模型后单独写推理脚本预测结果一团糟甚至报输入维度错误。原因训练脚本里 STFT 的窗长、步长、采样率和推理脚本里用的参数不完全一致。最常见的情况是训练时用的window_len256推理时为了追求更高频率分辨率改成了 1024输入张量的 shape 变了模型自然不认识。解决把 STFT 参数和数据预处理逻辑封装成同一个函数训练和推理都调用这个函数。参数一旦确定全流程不许改动。我在项目里的做法是把 STFT 参数写进配置文件训练脚本和推理脚本都从同一个配置文件读取。4.3 样本不均衡导致假高准确率现象总测试准确率 95%看起来很不错但看每个类别的准确率正常类 99%外圈故障 88%滚动体故障只有 40%。原因正常样本占了数据集的 80% 以上模型把几乎所有样本都判成正常类就能拿到很高的总体准确率少数类样本根本学不到特征。这类问题不看混淆矩阵很难发现。解决训练阶段使用加权随机采样器让每个类别的样本在训练中出现的概率大致相同或者使用类别权重交叉熵对样本少的类别加大误分类惩罚。评估阶段至少要看混淆矩阵和每个类别的 F1-score不能只看总准确率。4.4 GPU 和 CPU 推理结果不一致现象训练时 GPU 上验证集准确率 97%部署到 CPU 上变成 85%同一个模型权重。原因PyTorch 的批归一化层在训练和推理模式下行为不同。训练时用 batch 内统计量推理时用记录好的全局统计量。模型如果用了随机数据增强输入分布也会有差异。另外 CPU 和 GPU 的浮点计算精度差异也会造成微小偏差但通常不会导致准确率掉那么多。解决保存模型前确认model.eval()加载后先跑一遍测试集确认指标对得上再部署。如果 CPU 推理准确率仍然掉检查数据预处理流程尤其是归一化是否和训练时完全一致。4.5 图像尺寸不匹配现象训练代码里数据加载时做了 Resize 到 224但推理时忘记 Resize模型输入层报维度不匹配或者用尺寸大小不一的图像直接输入。原因训练和推理的数据流不一致其中一条路径漏了预处理。解决把完整的预处理管线包括 Resize、灰度转 RGB、Normalize统一封装成一个预处理函数训练、验证、推理全部复用。这样只要函数定义不动数据流就永远一致。5. 把模型做成诊断平台在线推理脚本与简易界面落地5.1 推理流程设计先判有无故障再判故障类型训练好模型只是第一步真正落地到平台时会发现一个工程问题:现场设备正常运行的时间远大于故障时间如果把每一帧数据都送进模型不仅浪费算力而且模型对正常样本的误报会让维护人员很快失去信任。我一般会在模型前面加一道门槛先用 RMS 均方根值判断振动能量是否超过阈值。轴承正常运行时振动能量在一个稳定区间一旦出现剥落或点蚀振动能量会明显抬升。RMS 低于阈值的直接判定正常只有 RMS 超标的片段才送进深度学习模型做细分类。这个做法能过滤掉 70% 以上的正常数据大大降低推理负载。5.2 流式数据推理脚本import torch import numpy as np from scipy.signal import stft fs 12800 window_len 256 hop_len 128 window_size fs # 每次滑动取 1 秒数据 stride fs // 2 # 0.5 秒滑动一次保证重叠覆盖 rms_threshold 2.5 # RMS 阈值现场根据实际工况标定 model load_model(best_model.pt, num_classes4) model.eval() device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) def compute_rms(data): return float(np.sqrt(np.mean(data ** 2))) def signal_to_spec_db(signal): _, _, Zxx stft(signal, fsfs, npersegwindow_len, noverlaphop_len, windowhann) spec np.abs(Zxx) spec_db 20 * np.log10(spec / spec.max() 1e-12) # 归一化到 0-255 并转为张量 img (spec_db - spec_db.min()) / (spec_db.max() - spec_db.min()) img (img * 255).astype(np.uint8) return img while True: data read_from_sensor(window_size) # 阻塞读取 1 秒数据 rms compute_rms(data) if rms rms_threshold: # 振动能量正常不进入模型标记为“正常”状态 push_status(normal, confidence0.0, rmsrms) continue spec_img signal_to_spec_db(data) # 灰度图转 RGB尺寸对齐训练配置 tensor torch.from_numpy(spec_img).float().unsqueeze(0).unsqueeze(0) tensor tensor.repeat(1, 3, 1, 1) # 复制成 3 通道 tensor torch.nn.functional.interpolate(tensor, size(224, 224)) tensor (tensor / 255.0 - 0.485) / 0.229 # 归一化按训练集统计量 tensor tensor.to(device) with torch.no_grad(): outputs model(tensor) probs torch.softmax(outputs, dim1) label torch.argmax(probs, dim1).item() confidence probs[0, label].item() push_status(ffault_{label}, confidenceconfidence, rmsrms)这个推理脚本的核心思路是两个阶段串联先算 RMS 做粗筛查再进模型做精分类。RMS 阈值设 2.5 是经验值现场需要采集一段正常运行的数据取 RMS 均值的 1.5 到 2 倍作为阈值直接套默认值很容易误报。注意归一化部分写的是减均值除以标准差直接按训练时的标准化公式做运本文还有配套的精品资源点击获取
基于深度学习的轴承故障诊断：从STFT时频图到ResNet18模型训练与部署
简介基于深度学习的轴承故障诊断平台面向机械故障诊断与深度学习方向的毕业设计、课程设计开发者针对传统诊断高度依赖专业知识、效率低下、难以适应复杂工况等痛点提供一套可运行、可扩展的完整方案。压缩包共63个文件、大小14.85MB其中含11个Python脚本覆盖数据预处理、特征提取、模型训练、模型评估、故障诊断等关键环节10个MAT格式的轴承振动数据集用于模型训练与验证34张可视化结果图包含训练曲线、混淆矩阵、ROC曲线、精确率召回率曲线等方便直观对比效果。平台集成了CNN、LSTM、GRU、MCNN-LSTM等多种网络并引入YOLO进行实时故障检测同时提供了界面化的交互工具和标准化处理功能可帮助使用者快速跑通从原始信号到故障分类的完整链路。目录结构清晰、模块划分明确适合论文复现、课程展示或毕设参考。目前已有59人学习下载。1. 轴承故障诊断为什么值得用深度学习去做传统轴承故障诊断靠老师傅拿听音棒贴在轴承座上听经验老到的人能从异响里分辨出外圈剥落还是内圈点蚀但这份手艺很难复制判断标准也说不清道不明。深度学习出现以后这个场景被彻底改写了——振动信号转成时频图喂给卷积网络模型自己学习故障特征准确率能做到 98% 以上而且整套流程是可以复现、可以参数化的。这份《基于深度学习的轴承故障诊断平台》资源覆盖的就是从原始振动信号到最终故障类型输出的完整链路包括信号预处理、时频图生成、模型训练和平台部署。适合正在做毕业设计、课程设计的同学也适合想快速搭一套诊断 demo 的工程师下载下来照着跑一遍就能把原理和代码对上了。2. 先把信号变成图像STFT 时频图生成与样本集制作2.1 为什么不能把原始振动信号直接丢给卷积网络很多初学者拿到轴承振动数据后第一反应是把一维信号直接拼成二维矩阵输入 CNN。这个做法不算错但效果通常不理想。原因在于卷积核的平移共享特性更适合捕捉图像中的局部模式而原始振动信号的故障特征分散在时间轴和频率轴两个维度上时域波形里特征被背景噪声淹没直接堆成一维卷积网络需要非常深的网络才能学到有区分度的特征小数据集上根本训不动。工程上的常见做法是先做时频变换把一维信号变成二维时频图。时频图横轴是时间纵轴是频率颜色深浅代表能量强度轴承故障的周期性冲击会表现为特定频带上的能量聚集这个模式在二维图像里非常明显可以直接套用成熟的图像分类网络。常见选择是短时傅里叶变换 STFT计算量可控物理含义直观配合 Hann 窗和 50% 重叠率能拿到不错的频率分辨率。比小波变换好调参比 Wigner-Ville 分布抗交叉项干扰作为入门方案性价比最高。2.2 Python 实现 STFT 时频图批量生成import numpy as np from scipy.signal import stft from PIL import Image import os fs 12800 # 采样率单位 Hz常见工业采集卡的默认配置 window_len 256 # 窗长决定频率分辨率 hop_len 128 # 步长与窗长配合实现 50% 重叠 total_nums 300 # 待处理的原始信号文件数量 output_dir stft_images os.makedirs(output_dir, exist_okTrue) for idx in range(total_nums): # 假设 read_vibration_signal 读入一段长度为 fs 秒的振动信号 signal read_vibration_signal(idx) f, t, Zxx stft(signal, fsfs, npersegwindow_len, noverlaphop_len, windowhann) # Zxx 是复数矩阵取幅值后转分贝 spec np.abs(Zxx) spec_db 20 * np.log10(spec / spec.max() 1e-12) # 分贝值归一化到 0-255保存为灰度图 img_arr (spec_db - spec_db.min()) / (spec_db.max() - spec_db.min()) img_arr (img_arr * 255).astype(np.uint8) Image.fromarray(img_arr).save( os.path.join(output_dir, fsample_{idx:04d}.png) )这段代码的核心逻辑是把每段原始振动信号做 STFT得到复数矩阵后取幅值并转换成分贝刻度。之所以用分贝而不是线性幅值是因为轴承故障冲击产生的能量和正常振动能量的差距往往有几个数量级线性刻度下微弱故障特征会被淹没分贝刻度能把这些细微差别拉开。窗长 256 对应的频率分辨率是 50 Hz对轴承故障特征频率来说这个数值是足够的。如果你的数据采集频率更高比如 48 kHz建议把窗长提到 512 或 1024否则每个频率仓覆盖范围太宽特征会糊在一起。步长 128 配合窗长 256 实现了 50% 重叠这是 STFT 的常见配置时间分辨率足够计算量也可接受。2.3 样本集制作与数据增强的取舍时频图生成只是第一步真正决定模型上限的是样本集的构造方式。一个常见误区是每段原始信号只生成一张图整个数据集稀疏得很模型很容易过拟合。我的做法是把长信号切成固定长度的小段每段时长 1 秒重叠 50%这样一段 10 秒的原始信号能切出 19 段数据集直接扩到将近 20 倍。数据增强方面要克制。对时频图做随机翻转、随机裁剪这类图像增强技巧在自然图像上很有效但放在时频图上需要谨慎。时间轴方向的小幅平移可以接受因为故障冲击本来就有随机性但频率轴的翻转会破坏频率和故障类型的对应关系——内圈故障和外圈故障的特征频率不同翻转之后模型学到的是错误映射。经验做法是只做两类增强给原始信号叠加高斯白噪声后再生成时频图模拟不同工况下的背景噪声变化以及时间轴上的小幅裁剪。这两类增强都在信号层面做生成图像后不再做任何图像级增强。另外每类故障的样本数要尽量均衡不然训练出来的模型会对样本多的类别过度偏好这个在后面的避坑章节会细说。3. 模型选型与训练流程从 ResNet18 微调到故障分类3.1 一维 CNN 还是二维 CNN先看你的数据形态选模型之前先想清楚一个问题你手上最终的数据形态是什么。如果你把信号转成了时频图那就用二维 CNN这是主流选择因为可以直接加载 ImageNet 预训练权重做迁移学习小数据集也能收敛得很稳。如果你坚持用原始一维信号那就得走 1D-CNN 路线网络结构得自己设计卷积核尺寸、空洞率、池化策略都要反复试调试成本高出一截。这份资源里默认的路线是二维 CNN 加 ResNet18 迁移学习。选 ResNet18 而不是 ResNet50 或更深的网络是因为轴承故障诊断的数据集规模通常不大深度网络的参数量在这种数据量下容易过拟合ResNet18 的残差结构已经足够提取时频图上的纹理特征。而且 ResNet18 在 CPU 上推理速度快后续做成实时诊断平台时不会成为性能瓶颈。如果你的数据量确实很大比如每类故障样本超过一万张可以考虑升级到 ResNet50。但训练轮数要相应增加学习率要调低否则深层网络在小数据集上很容易训出训练集准确率极高、验证集一塌糊涂的结果。3.2 PyTorch 迁移学习训练脚本及参数清单import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from torchvision import models, transforms from PIL import Image num_classes 4 # 正常、外圈故障、内圈故障、滚动体故障 batch_size 32 lr 1e-4 epochs 50 device torch.device(cuda if torch.cuda.is_available() else cpu) class FaultDataset(Dataset): def __init__(self, img_paths, labels, transformNone): self.img_paths img_paths self.labels labels self.transform transform def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img Image.open(self.img_paths[idx]).convert(L) # 灰度图转三通道适配 ImageNet 预训练权重 img img.convert(RGB) if self.transform: img self.transform(img) label self.labels[idx] return img, label train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 加载预训练 ResNet18替换最后全连接层 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, num_classes) model model.to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lrlr, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxepochs) train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workers4) for epoch in range(epochs): model.train() running_loss 0.0 correct 0 total 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() running_loss loss.item() scheduler.step() train_acc 100.0 * correct / total print(fEpoch [{epoch 1}/{epochs}] Loss: {running_loss / len(train_loader):.4f} fTrain Acc: {train_acc:.2f}%)代码逻辑分几块数据集类负责读图、转灰度再转三通道这是为了适配 ImageNet 预训练权重的输入要求模型部分替换了 ResNet18 的最后一层全连接输出类别数改成你的故障类别数训练循环是标准的 PyTorch 流程每轮训练完更新学习率。迁移学习的处理方式值得注意默认情况下所有层的权重都会参与训练这在数据量较大时没问题。如果数据量很少比如几百张图建议冻结前面的卷积层只训练最后一层全连接代码里加一行model.requires_grad_(False)再把model.fc.requires_grad_(True)训练速度更快过拟合风险也更低。训练参数上AdamW 配合 1e-4 初始学习率算是个稳健组合weight_decay 设 1e-4 控制权重范数。余弦退火学习率调度器比固定学习率好前期大步探索后期小步收敛不容易出现 loss 震荡。3.3 训练过程的监控指标与模型保存策略训练过程不能只看训练集准确率我最关心的是验证集准确率和 loss 曲线的间距。训练集 loss 持续下降但验证集 loss 不降反升这是典型的过拟合信号。常见的应对手段有三个加大数据增强的强度、增加 dropout、降低模型容量。对迁移学习场景来说降低模型容量换成 ResNet18 就已经是低配了优先考虑前两种。模型保存有个血泪经验不要每轮都保存磁盘很快就满了。我的习惯是保存验证集准确率最高的那一轮同时只保留最近五个检查点代码里用一个列表记录准确率每次有新高就覆盖旧的。另外保存时把模型结构定义和权重一起打包用torch.save(model.state_dict(), path)就够了加载时先定义好网络结构再load_state_dict。训练完成后再拿测试集测一次泛化指标注意测试集不能参与过任何训练和验证流程否则指标没有参考意义。4. 避坑与常见问题排查训练跑崩和假高准确率的五个根源4.1 模型准确率虚高归一化统计量泄漏现象训练集准确率 99%测试集准确率也 99%但部署到现场数据上准确率掉到 60% 以下。原因对整批数据统一计算均值和标准差做归一化或者统一做了最大值最小值缩放到 0-255这会让归一化统计量同时包含训练集和测试集的信息相当于测试集的信息提前泄露给了模型训练过程。测试时新数据按训练集的统计量归一化分布对不上准确率自然崩。解决归一化统计量只能在训练集上计算然后把训练集的均值和标准差保存到文件测试时加载同一个文件使用。代码上就是transform里 Normalize 的 mean 和 std 用训练集算出来的值测试集和推理阶段用同样的值.4.2 时频图参数不一致现象训练时准确率很高保存模型后单独写推理脚本预测结果一团糟甚至报输入维度错误。原因训练脚本里 STFT 的窗长、步长、采样率和推理脚本里用的参数不完全一致。最常见的情况是训练时用的window_len256推理时为了追求更高频率分辨率改成了 1024输入张量的 shape 变了模型自然不认识。解决把 STFT 参数和数据预处理逻辑封装成同一个函数训练和推理都调用这个函数。参数一旦确定全流程不许改动。我在项目里的做法是把 STFT 参数写进配置文件训练脚本和推理脚本都从同一个配置文件读取。4.3 样本不均衡导致假高准确率现象总测试准确率 95%看起来很不错但看每个类别的准确率正常类 99%外圈故障 88%滚动体故障只有 40%。原因正常样本占了数据集的 80% 以上模型把几乎所有样本都判成正常类就能拿到很高的总体准确率少数类样本根本学不到特征。这类问题不看混淆矩阵很难发现。解决训练阶段使用加权随机采样器让每个类别的样本在训练中出现的概率大致相同或者使用类别权重交叉熵对样本少的类别加大误分类惩罚。评估阶段至少要看混淆矩阵和每个类别的 F1-score不能只看总准确率。4.4 GPU 和 CPU 推理结果不一致现象训练时 GPU 上验证集准确率 97%部署到 CPU 上变成 85%同一个模型权重。原因PyTorch 的批归一化层在训练和推理模式下行为不同。训练时用 batch 内统计量推理时用记录好的全局统计量。模型如果用了随机数据增强输入分布也会有差异。另外 CPU 和 GPU 的浮点计算精度差异也会造成微小偏差但通常不会导致准确率掉那么多。解决保存模型前确认model.eval()加载后先跑一遍测试集确认指标对得上再部署。如果 CPU 推理准确率仍然掉检查数据预处理流程尤其是归一化是否和训练时完全一致。4.5 图像尺寸不匹配现象训练代码里数据加载时做了 Resize 到 224但推理时忘记 Resize模型输入层报维度不匹配或者用尺寸大小不一的图像直接输入。原因训练和推理的数据流不一致其中一条路径漏了预处理。解决把完整的预处理管线包括 Resize、灰度转 RGB、Normalize统一封装成一个预处理函数训练、验证、推理全部复用。这样只要函数定义不动数据流就永远一致。5. 把模型做成诊断平台在线推理脚本与简易界面落地5.1 推理流程设计先判有无故障再判故障类型训练好模型只是第一步真正落地到平台时会发现一个工程问题:现场设备正常运行的时间远大于故障时间如果把每一帧数据都送进模型不仅浪费算力而且模型对正常样本的误报会让维护人员很快失去信任。我一般会在模型前面加一道门槛先用 RMS 均方根值判断振动能量是否超过阈值。轴承正常运行时振动能量在一个稳定区间一旦出现剥落或点蚀振动能量会明显抬升。RMS 低于阈值的直接判定正常只有 RMS 超标的片段才送进深度学习模型做细分类。这个做法能过滤掉 70% 以上的正常数据大大降低推理负载。5.2 流式数据推理脚本import torch import numpy as np from scipy.signal import stft fs 12800 window_len 256 hop_len 128 window_size fs # 每次滑动取 1 秒数据 stride fs // 2 # 0.5 秒滑动一次保证重叠覆盖 rms_threshold 2.5 # RMS 阈值现场根据实际工况标定 model load_model(best_model.pt, num_classes4) model.eval() device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) def compute_rms(data): return float(np.sqrt(np.mean(data ** 2))) def signal_to_spec_db(signal): _, _, Zxx stft(signal, fsfs, npersegwindow_len, noverlaphop_len, windowhann) spec np.abs(Zxx) spec_db 20 * np.log10(spec / spec.max() 1e-12) # 归一化到 0-255 并转为张量 img (spec_db - spec_db.min()) / (spec_db.max() - spec_db.min()) img (img * 255).astype(np.uint8) return img while True: data read_from_sensor(window_size) # 阻塞读取 1 秒数据 rms compute_rms(data) if rms rms_threshold: # 振动能量正常不进入模型标记为“正常”状态 push_status(normal, confidence0.0, rmsrms) continue spec_img signal_to_spec_db(data) # 灰度图转 RGB尺寸对齐训练配置 tensor torch.from_numpy(spec_img).float().unsqueeze(0).unsqueeze(0) tensor tensor.repeat(1, 3, 1, 1) # 复制成 3 通道 tensor torch.nn.functional.interpolate(tensor, size(224, 224)) tensor (tensor / 255.0 - 0.485) / 0.229 # 归一化按训练集统计量 tensor tensor.to(device) with torch.no_grad(): outputs model(tensor) probs torch.softmax(outputs, dim1) label torch.argmax(probs, dim1).item() confidence probs[0, label].item() push_status(ffault_{label}, confidenceconfidence, rmsrms)这个推理脚本的核心思路是两个阶段串联先算 RMS 做粗筛查再进模型做精分类。RMS 阈值设 2.5 是经验值现场需要采集一段正常运行的数据取 RMS 均值的 1.5 到 2 倍作为阈值直接套默认值很容易误报。注意归一化部分写的是减均值除以标准差直接按训练时的标准化公式做运本文还有配套的精品资源点击获取
