轴承故障诊断实战：小波时频图与Swin Transformer端到端方案-北京尧图网络

简介这份资源面向具备Python与深度学习基础的科研人员、研究生及工业设备诊断工程师提供一套基于小波时频图WTFP结合移位窗口视觉TransformerST的轴承故障诊断完整项目实例。它解决非平稳振动、工况变化与跨设备迁移下传统方法诊断性能不佳的问题覆盖风力发电、工业电机、数控机床、轨道交通等旋转机械场景。压缩包共1个docx文件约147KB以文档形式系统呈现从原始信号切片、标准化、连续小波变换生成时频图、数据增强、模型训练到测试评估的端到端链路并包含GUI设计与代码详解。已有63人学习。读者可获取复Morlet小波多分辨率特征提取、窗口多头自注意力与移位窗口跨区域交互、Patch Merging及分类头等关键模块的实现思路以及数据划分、参数调优与训练稳定性设计为复现高效鲁棒的工业智能运维原型提供可参考的工程方案。1. 轴承故障诊断为什么绕不开小波时频图与移位窗口 Transformer产线上最让人头疼的往往不是电机彻底烧掉而是轴承早期那点几乎听不见的异响。加速度传感器采回来的原始振动信号在时域上就是一条上下乱跳的曲线肉眼看不出任何规律直接丢给普通分类网络准确率经常在及格线附近晃。原因不复杂轴承故障的特征频率外圈、内圈、滚动体、保持架藏在冲击响应的调制里时域看不出纯频域又丢掉了冲击发生的时刻。小波时频图Wavelet Time-Frequency Plotwtfp正好把这两个维度同时保留下来把一维振动信号变成一张二维图像故障特征在图上表现为周期性的亮条纹。而移位窗口视觉 TransformerSwin Transformer常简写 ST擅长在图像上做多尺度、带局部窗口注意力的特征提取比传统 CNN 更能抓住这些条纹的全局周期结构。这套组合在近两年的轴承故障诊断里被反复验证适合手里有振动数据、想从「手工特征浅层分类器」升级到端到端深度模型的工程师。下面我按自己落地的顺序把数据、时频图生成、模型、GUI 和踩坑一次讲透。2. 从振动信号到小波时频图数据准备与图像生成2.1 为什么选连续小波变换而不是短时傅里叶短时傅里叶变换STFT的窗长一旦定死时间分辨率和频率分辨率就固定了而轴承冲击既需要高频段的时间精度又需要低频段的频率精度STFT 很难两头兼顾。连续小波变换CWT用可伸缩的母小波高频处窗口窄、低频处窗口宽天然适配冲击类信号。工程上我一般选复 Morlet 小波它在时频域都集中生成的时频图条纹清晰、背景干净比墨西哥帽小波更适合做图像输入。选型确定后参数就三个尺度范围、采样频率、母小波类型。尺度范围决定频率覆盖采样频率决定奈奎斯特上限母小波决定时频聚集性。这三个参数不调好后面模型再强也白搭。2.2 用 PyWavelets 生成时频图的最小脚本下面这段代码把一条一维振动信号转成归一化后的时频图直接可以喂给后面的 Swin Transformer。数据集我用的是常见的 CWRU 轴承数据格式这里只讲格式不涉及任何具体机构每条记录是一段驱动端加速度信号。import numpy as np import pywt import matplotlib.pyplot as plt from PIL import Image def signal_to_wtfp(signal, fs12000, scalesNone, waveletcmor1.5-1.0, img_size(224, 224), save_pathNone): signal: 一维振动信号 (numpy array) fs: 采样频率CWRU 驱动端常用 12kHz scales: 小波尺度序列默认覆盖 10Hz~fs/2 wavelet: 复 Morlet带宽-中心频率比 1.5-1.0 img_size: 输出图像尺寸Swin 常用 224x224 if scales is None: # 尺度与频率近似关系: f ≈ fc * fs / scalefc 取 1.0 scales np.arange(1, 128) # 连续小波变换返回复数系数 coeffs, freqs pywt.cwt(signal, scales, wavelet, sampling_period1.0/fs) # 取模得到幅值时频图转 dB 压缩动态范围 magnitude np.abs(coeffs) magnitude 20 * np.log10(magnitude 1e-12) # 归一化到 0~255方便存成灰度图 mag_min, mag_max magnitude.min(), magnitude.max() norm (magnitude - mag_min) / (mag_max - mag_min 1e-12) img (norm * 255).astype(np.uint8) # 缩放到模型输入尺寸 img_pil Image.fromarray(img).resize(img_size, Image.BILINEAR) if save_path: img_pil.save(save_path) return np.array(img_pil) # 示例读取一段信号并生成时频图 # signal np.loadtxt(bearing_signal.txt) # 一维数组 # wtfp signal_to_wtfp(signal, fs12000, save_pathwtfp_sample.png)逻辑说明pywt.cwt返回的是复数系数矩阵行对应尺度、列对应时间点取模后就是时频幅值图。转 dB 是为了压缩动态范围否则强冲击会把弱特征淹没。归一化到 0~255 是为了存成标准图像格式方便后续用图像增强。参数方面scales从 1 到 127 覆盖了大约 94Hz 到 12kHz基本覆盖轴承故障特征频段如果你的采样频率是 48kHz尺度上限要相应放大到 512 左右否则高频段会被截断。waveletcmor1.5-1.0里的 1.5 是带宽参数1.0 是中心频率带宽越大时频图越平滑但频率分辨率下降我一般从 1.5 起步条纹模糊就降到 1.0。2.3 数据集划分与标签映射时频图生成后按故障类型和损伤直径建文件夹目录结构直接决定标签。常见做法是每个类别取 80% 做训练、20% 做测试且必须按「时间段」划分而不是随机打乱——同一段连续信号切出来的相邻样本高度相关随机划分会让测试集泄漏训练集信息准确率虚高到 99% 但换一台设备就崩。我一般按采集批次划分同一批次的数据只出现在训练集或测试集之一。标签映射用一个简单的字典维护类别数取决于你的数据CWRU 常见是 10 类正常 3 种故障位置 × 3 种损伤直径。类别不平衡时用加权采样别直接上过采样时频图过采样容易生成不真实的条纹。3. Swin Transformer 做时频图分类结构改造与训练配置3.1 为什么 Swin 比普通 ViT 更适合时频图普通 ViT 把图像切成固定 patch 后做全局自注意力计算量随分辨率平方增长而且时频图里真正有用的信息集中在少数亮条纹上全局注意力会被大量背景噪声稀释。Swin 用移位窗口把注意力限制在局部窗口内再通过窗口移位实现跨窗口信息交换计算量线性增长同时保留了多尺度层次结构。时频图的故障条纹是局部纹理加全局周期Swin 的层次化设计正好匹配浅层抓局部条纹边缘深层抓周期重复模式。3.2 用 timm 加载 Swin 并改造分类头下面用timm加载预训练 Swin-Tiny把输入改成单通道灰度输出改成你的类别数。预训练权重在 ImageNet 上学的纹理特征对时频图有迁移价值比从头训练收敛快很多。import torch import torch.nn as nn import timm class SwinWTFPClassifier(nn.Module): def __init__(self, num_classes10, in_chans1, pretrainedTrue): super().__init__() # 加载 Swin-Tiny修改输入通道和分类数 self.backbone timm.create_model( swin_tiny_patch4_window7_224, pretrainedpretrained, in_chansin_chans, # 灰度时频图单通道 num_classesnum_classes ) def forward(self, x): return self.backbone(x) # 训练配置 def build_optimizer(model, lr1e-4, weight_decay0.05): # 分层学习率backbone 小分类头大 params [ {params: model.backbone.parameters(), lr: lr}, ] return torch.optim.AdamW(params, lrlr, weight_decayweight_decay) # 示例 # model SwinWTFPClassifier(num_classes10) # optimizer build_optimizer(model) # criterion nn.CrossEntropyLoss(label_smoothing0.1)逻辑说明in_chans1让 timm 自动把第一层卷积的权重在通道维度求平均适配灰度输入这是迁移学习里最省事的做法。num_classes直接替换分类头。优化器用 AdamW权重衰减 0.05 是 Swin 的常用值太大欠拟合、太小过拟合。学习率 1e-4 是微调预训练模型的稳妥起点如果你数据量少于 2000 张降到 5e-5 更稳。label_smoothing0.1对类别不平衡和标注噪声有抑制作用轴承数据里偶尔有标错的样本这个参数能救回来一点。3.3 训练循环与关键超参训练循环本身不复杂关键是几个超参batch size、学习率调度、混合精度。Swin-Tiny 在 224×224 下单卡 8GB 显存能跑 batch size 16用混合精度可以到 32。学习率用余弦退火warmup 5 个 epoch避免一开始就把预训练权重冲垮。from torch.cuda.amp import autocast, GradScaler from torch.optim.lr_scheduler import CosineAnnealingLR def train_one_epoch(model, loader, optimizer, criterion, scaler, device): model.train() total_loss, correct, total 0.0, 0, 0 for imgs, labels in loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() with autocast(): # 混合精度省显存提速 outputs model(imgs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() total_loss loss.item() * imgs.size(0) correct (outputs.argmax(1) labels).sum().item() total imgs.size(0) return total_loss / total, correct / total # 训练主循环骨架 # scaler GradScaler() # scheduler CosineAnnealingLR(optimizer, T_max50, eta_min1e-6) # for epoch in range(50): # loss, acc train_one_epoch(model, train_loader, optimizer, criterion, scaler, device) # scheduler.step()逻辑说明autocast把大部分运算降到 FP16显存占用减半速度提升约 30%但要注意 loss 计算和梯度更新仍在 FP32GradScaler负责缩放梯度防止下溢。余弦退火让学习率从 1e-4 平滑降到 1e-6后期微调更精细。50 个 epoch 对 Swin 微调通常够用如果验证集准确率在 20 个 epoch 后还在涨可以加到 80。早停策略用验证集准确率连续 10 个 epoch 不提升就停别硬跑。4. GUI 设计与推理部署让模型能被人用起来4.1 用 PyQt5 搭一个最小可用界面模型训完只是半成品产线上的人不会敲命令行。GUI 我一般用 PyQt5功能就三块选信号文件、显示时频图、给出诊断结果和置信度。界面不用花哨能跑通就行。import sys import numpy as np from PyQt5.QtWidgets import (QApplication, QMainWindow, QPushButton, QLabel, QVBoxLayout, QWidget, QFileDialog) from PyQt5.QtGui import QPixmap import torch from PIL import Image class DiagnosisGUI(QMainWindow): def __init__(self, model, device): super().__init__() self.model model self.device device self.setWindowTitle(轴承故障诊断) self.resize(500, 600) self.btn_load QPushButton(选择振动信号文件) self.btn_load.clicked.connect(self.load_signal) self.lbl_img QLabel(时频图将显示在这里) self.lbl_result QLabel(诊断结果待检测) layout QVBoxLayout() layout.addWidget(self.btn_load) layout.addWidget(self.lbl_img) layout.addWidget(self.lbl_result) container QWidget() container.setLayout(layout) self.setCentralWidget(container) def load_signal(self): path, _ QFileDialog.getOpenFileName(self, 选择信号, , Text (*.txt *.csv)) if not path: return signal np.loadtxt(path) # 复用第 2 章的时频图生成函数 from wtfp_utils import signal_to_wtfp img signal_to_wtfp(signal, fs12000, save_pathtmp_wtfp.png) pixmap QPixmap(tmp_wtfp.png).scaled(300, 300) self.lbl_img.setPixmap(pixmap) # 推理 tensor torch.from_numpy(img).float().unsqueeze(0).unsqueeze(0) / 255.0 tensor tensor.to(self.device) self.model.eval() with torch.no_grad(): logits self.model(tensor) prob torch.softmax(logits, dim1) conf, pred prob.max(1) self.lbl_result.setText(f诊断结果类别 {pred.item()}置信度 {conf.item():.3f}) # if __name__ __main__: # app QApplication(sys.argv) # model SwinWTFPClassifier(num_classes10).to(cuda).eval() # model.load_state_dict(torch.load(best_swin.pth)) # gui DiagnosisGUI(model, cuda) # gui.show() # sys.exit(app.exec_())逻辑说明load_signal里读文件、生成时频图、推理三步串起来。推理时输入要归一化到 0~1和训练时的预处理保持一致否则准确率会掉。unsqueeze(0)两次分别加 batch 维和通道维。置信度用 softmax 最大值低于 0.6 时建议在界面上标黄提示「结果存疑」产线上误报比漏报更烦人。4.2 推理速度与显存优化Swin-Tiny 在单张 224×224 灰度图上推理约 8msRTX 3060完全够实时。如果部署在 CPU 上用torch.jit.trace导出 TorchScript再开torch.set_num_threads(4)单张约 60ms也能接受。显存不够就把模型转成 FP16model.half()输入也转 half显存减半精度损失通常小于 0.5%。批量推理时注意时频图生成是 CPU 瓶颈pywt.cwt对长信号比较慢。我一般先把所有测试信号预生成时频图存成 PNG推理时只读图避免每次现场算 CWT。5. 避坑与排查这套方案最容易翻车的五个地方5.1 时频图全黑或全白现象生成的时频图要么一片黑要么一片白看不出任何条纹。原因通常是 dB 转换后动态范围没压好或者信号幅值太小。解决先检查原始信号幅值范围如果都在 1e-6 量级先乘 1e6 再算 CWTdB 转换后手动设vmin和vmax比如取 5% 和 95% 分位数别用 min/max异常值会把整张图拉偏。5.2 测试集准确率 99% 但现场一塌糊涂现象离线测试准确率极高换一批数据就崩。原因几乎都是数据泄漏——同一段连续信号切出的相邻样本被随机分到了训练和测试集。解决按采集批次或时间段划分同一批次只进一个集合。另外检查时频图生成时有没有用到全局统计量比如全局归一化如果有改成按单样本归一化。5.3 训练 loss 不降或震荡现象loss 在前几个 epoch 就卡住或者上下剧烈震荡。原因可能是学习率太大把预训练权重冲垮了或者 batch size 太小导致梯度噪声大。解决学习率降到 5e-5warmup 加到 10 个 epochbatch size 至少 16不够就累积梯度。另外检查输入归一化时频图如果没除 255输入值在 0~255 范围第一层就爆了。5.4 类别不平衡导致少数类全错现象正常样本识别很准但某种损伤直径的样本几乎全错。原因是该类样本太少模型偏向多数类。解决用加权 CrossEntropyLoss权重取类别频率的倒数或者用 WeightedRandomSampler 在 DataLoader 层面重采样。别用图像过采样时频图过采样生成的副本没有新信息反而加剧过拟合。5.5 GUI 推理结果和训练时对不上现象GUI 里显示的诊断结果和离线测试脚本不一致。原因通常是预处理不一致——训练时用了某种归一化或 resize 插值方式GUI 里用了另一种。解决把预处理封装成一个函数训练和推理都调同一个别在两处各写一遍。另外检查模型是否忘了.eval()BatchNorm 和 Dropout 在训练模式下会改变输出。6. 进阶技巧用 Grad-CAM 看模型到底在看哪里模型给出结果只是第一步产线上的人会问「你凭什么说是外圈故障」。这时候 Grad-CAM 就派上用场了它能把模型注意力热力图叠加回时频图让你看到模型关注的是哪段频率、哪个时刻。对 Swin 这种层次化结构取最后一个 stage 的特征图做 CAM 效果最好。import torch import numpy as np import cv2 def swin_gradcam(model, input_tensor, target_layer, class_idxNone): model: SwinWTFPClassifier input_tensor: (1, 1, 224, 224) target_layer: 最后一个 stage 的 norm 层 features, grads [], [] def forward_hook(module, inp, out): features.append(out) def backward_hook(module, grad_in, grad_out): grads.append(grad_out[0]) handle_f target_layer.register_forward_hook(forward_hook) handle_b target_layer.register_full_backward_hook(backward_hook) model.eval() output model(input_tensor) if class_idx is None: class_idx output.argmax(1).item() model.zero_grad() output[0, class_idx].backward() feat features[0].detach() # (1, C, H, W) grad grads[0].detach() # (1, C, H, W) weights grad.mean(dim(2, 3), keepdimTrue) # 全局平均池化 cam (weights * feat).sum(dim1, keepdimTrue) cam torch.relu(cam) cam cam.squeeze().cpu().numpy() cam (cam - cam.min()) / (cam.max() - cam.min() 1e-8) cam cv2.resize(cam, (224, 224)) handle_f.remove() handle_b.remove() return cam # 叠加显示 # cam swin_gradcam(model, tensor, model.backbone.layers[-1].blocks[-1].norm1) # heatmap cv2.applyColorMap((cam * 255).astype(np.uint8), cv2.COLORMAP_JET) # overlay cv2.addWeighted(orig_img_rgb, 0.6, heatmap, 0.4, 0)逻辑说明register_forward_hook抓特征图register_full_backward_hook抓梯度两者相乘再全局平均池化就得到 CAM。target_layer选最后一个 stage 的 norm 层太浅的层感受野小热力图会碎太深的层分辨率低定位不准。torch.relu只保留对目标类有正贡献的区域负贡献直接置零。叠加时原图权重 0.6、热力图 0.4既能看清条纹又能看到关注区域。验证方法上我一般做两件事一是看热力图是否落在故障特征频率对应的尺度带上如果模型关注的是低频背景说明它没学到真东西二是做遮挡测试把热力图高亮区域遮掉再推理准确率应该明显下降否则说明模型靠的是别处的伪特征。这两个测试比单纯看准确率靠谱得多。我自己的习惯是每训完一个模型先跑 20 张测试图的 Grad-CAM肉眼过一遍。有次发现模型全在看时频图边缘的空白区查了半天是生成时频图时 resize 引入了边框伪影去掉边框后准确率反而涨了两个点。这种问题光看 loss 曲线永远发现不了。希望帮到你。本文还有配套的精品资源点击获取

轴承故障诊断实战：小波时频图与Swin Transformer端到端方案

相关推荐

基于Java的教务系统开发：从表结构设计到并发选课的完整实战

if语句的执行规则：从真值表到短路求值，你真的理解了吗？

C++与FPGA协同设计实战：从架构划分到频谱仪软硬联调

更多资讯

pstack-claude实战：用AI辅助分析进程栈与线上排障

如何打造无可挑剔的代码质量检查工具：从需求到落地的工程实践

Windows 上跑 Codex 总卡第一步？Node.js 与 npm 环境配置避坑指南

运维和网工哪个发展好？从日常、技能栈到发展路径的全面对比

SQL Server病房管理系统课程设计：从E-R图到建表避坑指南

t3code 实战：构建本地化代码质量分析与复杂度度量体系

今日更新

PHP htmlentities()函数用法讲解

AI Agent工程实战：从七要素到七个决策点的系统设计指南

AI入局芯片设计验证：验证收敛50倍加速的真相与应对

本周精选

CMOS设计三大物理效应：体效应、沟道长度调制与亚阈值导电

微信小程序在线阅读毕设全解析：Java后端+MySQL+避坑实战

Android音频问题五层分析法：从App到Codec的系统化诊断

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？