☰
轴承故障诊断新范式:小波时频图与移位窗口Transformer实战
2026/10/1 14:22:17 网站建设 项目流程

简介:这份资源面向具备Python与深度学习基础的科研人员、研究生及工业设备状态监测工程师,提供一套基于小波时频图(WTFP)结合移位窗口视觉Transformer(ST)的轴承故障诊断完整项目实例。它针对非平稳振动、工况变化与跨设备迁移下传统方法诊断性能不佳的问题,构建了从原始信号切片、标准化、连续小波变换生成时频图、数据增强、模型训练到测试评估的端到端链路,并覆盖GPU推理、服务接口与监控维护等部署环节。资源包为1个docx文档,约147KB,内含完整程序代码、GUI设计与逐段代码详解,目录涵盖项目背景、模型架构、窗口注意力与移位窗口模块、训练验证及测试评估等章节。已有63人学习。读者可据此复现可运行的故障识别原型,掌握小波参数选择、数据划分与训练稳定性等工程要点,为风力发电、工业电机、数控机床等旋转机械的智能运维提供可落地的技术参考。

1. 轴承故障诊断为什么值得用「小波时频图 + 移位窗口 Transformer」重做一遍

振动信号是旋转机械最诚实的"体检报告",但这份报告是时变的、非平稳的,轴承一旦出现早期点蚀或剥落,冲击成分往往埋在强噪声和工频干扰里。传统做法要么靠人工特征(峭度、包络谱、均方根)加浅层分类器,要么直接把一维时序丢进 CNN,结果在变工况、变转速下泛化能力很差。这几年我反复试过一条路线:先把一维振动信号转成小波时频图(WTFP),把时间-频率的二维纹理显式暴露出来,再用移位窗口视觉 Transformer(ST,Swin Transformer 那一类结构)去做分类。它解决的核心问题是——让模型既能看到局部冲击的细节,又能建模全局的频带依赖,同时把计算量压到工程可接受的范围。这套方案适合做设备状态监测、故障诊断方向的工程师和研究生,尤其是手里已经有 CWRU 这类公开轴承数据集、想快速跑通一条完整链路的人。下面我按"信号怎么变图、图怎么进网络、网络怎么落地成 GUI"的顺序,把能抄作业的部分全写出来。

2. 从一维振动到小波时频图:WTFP 的生成链路与参数选择

2.1 为什么选小波而不是 STFT

短时傅里叶变换(STFT)的窗长一旦定死,时间分辨率和频率分辨率就固定了,而轴承故障的冲击成分频率高、持续时间短,低频的转频调制又需要长窗。小波变换用可伸缩的母小波,高频处窗窄、低频处窗宽,天然适配这种"多尺度"需求。工程上我一般用连续小波变换(CWT)生成时频图,母小波优先选 Morlet 或复 Morlet(cmor),因为它在时频聚集性上表现稳定,做出来的图纹理清晰,喂给视觉模型时特征更"好认"。

选型上还有几个现实考量:一是 CWT 的尺度序列要和采样频率、关心的故障特征频率对齐,不能随便取;二是生成的时频图要归一化到 0-255 存成灰度或伪彩色图,方便后续按图像任务处理;三是数据量控制,CWT 计算量大,通常先对信号分段(比如每段 1024 或 2048 点),再逐段变换。

2.2 用 PyWavelets 生成时频图的完整代码

import numpy as np import pywt import matplotlib.pyplot as plt from PIL import Image def signal_to_wtfp(signal, fs=12000, scales=None, wavelet='cmor1.5-1.0', img_size=(224, 224), save_path=None): """ 将一维振动信号转成小波时频图 signal: 1D array, 单段振动信号 fs: 采样频率, CWRU 驱动端常用 12kHz scales: 尺度序列, None 则自动生成 wavelet: 复 Morlet, 带宽-中心频率比 1.5-1.0 img_size: 输出图像尺寸, 对齐 Swin 常用输入 """ if scales is None: # 尺度覆盖到关心的频带, 一般 1~128 够用 scales = np.arange(1, 129) # 连续小波变换, 得到复数系数 coeffs, freqs = pywt.cwt(signal, scales, wavelet, sampling_period=1/fs) # 取模得到幅值谱, 转 dB 增强弱特征 mag = np.abs(coeffs) mag = 20 * np.log10(mag + 1e-12) # 归一化到 0-255 mag = (mag - mag.min()) / (mag.max() - mag.min() + 1e-12) img = (mag * 255).astype(np.uint8) # 缩放到网络输入尺寸 img = np.array(Image.fromarray(img).resize(img_size, Image.BILINEAR)) if save_path: Image.fromarray(img).save(save_path) return img, freqs # 示例: 读取一段 CWRU 数据并生成图 # data = np.loadtxt('97.mat') # 实际按数据集格式读取 # seg = data[:2048] # img, freqs = signal_to_wtfp(seg, fs=12000, save_path='wtfp_sample.png')

这段代码的逻辑是:先做 CWT 拿到复数系数矩阵,取模后转 dB,再归一化并缩放。参数上有三个关键点。scales决定频率覆盖范围,尺度越大对应频率越低,如果你的故障特征频率集中在 2-5kHz,尺度序列要相应调整,别一股脑取到 128 导致低频区全是噪声。wavelet里的cmor1.5-1.0是带宽和中心频率的比值,1.5 偏时域聚集、1.0 偏频域,实测下来这个组合在轴承冲击上比较均衡。img_size建议和后面 Swin 的输入保持一致,224 是通用选择,但如果你显存紧张,128 也能跑,只是细纹理会有损失。

提示:CWT 对每段信号都要算一次,数据量大时用多进程或先把所有图离线生成好存成图片,训练时直接读图,别在Dataset里实时算,否则 GPU 会一直等 CPU。

2.3 数据集组织与标签映射

生成完图之后,目录结构建议按类别分文件夹,方便直接用ImageFolder或自定义 Dataset。CWRU 常见做法是按故障直径(0.007、0.014、0.021 英寸)和内圈/外圈/滚动体组合成 10 类左右。我一般会写一个预处理脚本,把每类信号切成固定长度段,逐段生成图并落到对应文件夹,同时记录一个label_map.json,避免后面标签对不上。

import os, json from pathlib import Path def build_dataset(raw_dir, out_dir, seg_len=2048, stride=1024): label_map = {} idx = 0 for cls_name in sorted(os.listdir(raw_dir)): cls_path = Path(raw_dir) / cls_name if not cls_path.is_dir(): continue label_map[cls_name] = idx save_dir = Path(out_dir) / cls_name save_dir.mkdir(parents=True, exist_ok=True) for f in cls_path.glob('*.mat'): # 这里按实际 mat 结构读取, 假设拿到一维数组 data # data = load_mat_signal(f) # for i, start in enumerate(range(0, len(data)-seg_len, stride)): # seg = data[start:start+seg_len] # signal_to_wtfp(seg, save_path=save_dir / f'{f.stem}_{i}.png') pass idx += 1 with open(Path(out_dir) / 'label_map.json', 'w') as fp: json.dump(label_map, fp, ensure_ascii=False, indent=2) return label_map

分段时seg_len和stride要配合:段太长,一段里可能混多种状态;段太短,CWT 的频率分辨率不够。2048 点、步长 1024 是我常用的折中,重叠一半能增加样本量,也避免边界信息丢失。标签映射一定要落盘,训练和推理共用同一份,否则 GUI 里显示的结果会张冠李戴。

3. 移位窗口视觉 Transformer 怎么接住时频图

3.1 Swin 的层级结构与窗口注意力为什么适合时频图

标准 ViT 把图切成固定 patch 后做全局注意力,计算量随分辨率平方增长,而且时频图里真正有判别力的往往只是几条冲击带,全局注意力容易把噪声也一起加权。Swin 的做法是分层下采样加窗口内注意力,窗口之间再通过移位(shifted window)建立跨窗连接。对时频图来说,局部窗口能聚焦某段频率上的冲击纹理,移位机制又能把不同频带、不同时间片的信息串起来,这正好对应"局部冲击 + 全局频带依赖"的需求。

另一个工程优势是计算量。窗口注意力把复杂度从全局的平方降到与窗口大小相关的线性级别,224 输入下普通显卡也能训。层级结构还会输出多尺度特征,后面接分类头或做可视化都方便。

3.2 用 timm 搭建 Swin 分类网络

import torch import torch.nn as nn import timm class BearingSwin(nn.Module): def __init__(self, num_classes=10, model_name='swin_tiny_patch4_window7_224', pretrained=True, drop_rate=0.2): super().__init__() # timm 直接给出现成 Swin, 改分类头即可 self.backbone = timm.create_model( model_name, pretrained=pretrained, num_classes=num_classes, drop_rate=drop_rate ) def forward(self, x): return self.backbone(x) def build_model(num_classes=10, device='cuda'): model = BearingSwin(num_classes=num_classes) model.to(device) return model # 快速验证输入输出 # model = build_model(num_classes=10, device='cpu') # x = torch.randn(2, 3, 224, 224) # print(model(x).shape) # torch.Size([2, 10])

这里用timm是为了少造轮子,swin_tiny_patch4_window7_224是参数量和精度比较平衡的一档,预训练权重能显著加快收敛,尤其在轴承数据样本有限时。drop_rate设 0.2 是防止过拟合的常规手段,如果你的数据集很小(每类几百张),可以加到 0.3 并配合数据增强。注意输入必须是 3 通道,而我们的 WTFP 是单通道灰度图,所以 Dataset 里要复制成三通道,或者改模型第一层卷积,前者更省事。

3.3 训练循环与关键超参

from torch.utils.data import DataLoader from torchvision import datasets, transforms import torch.optim as optim def train_one_epoch(model, loader, criterion, optimizer, 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() outputs = model(imgs) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() * imgs.size(0) correct += (outputs.argmax(1) == labels).sum().item() total += imgs.size(0) return total_loss / total, correct / total def build_loaders(data_dir, batch_size=32, img_size=224): tf = transforms.Compose([ transforms.Grayscale(num_output_channels=3), # 单通道转三通道 transforms.Resize((img_size, img_size)), transforms.RandomHorizontalFlip(p=0.3), # 时频图翻转要谨慎 transforms.ToTensor(), transforms.Normalize(mean=[0.5]*3, std=[0.5]*3) ]) ds = datasets.ImageFolder(data_dir, transform=tf) loader = DataLoader(ds, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True) return loader, ds.classes # 训练主流程 # device = 'cuda' if torch.cuda.is_available() else 'cpu' # model = build_model(num_classes=10, device=device) # loader, classes = build_loaders('wtfp_dataset') # criterion = nn.CrossEntropyLoss() # optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-2) # for epoch in range(50): # loss, acc = train_one_epoch(model, loader, criterion, optimizer, device) # print(f'epoch {epoch}, loss {loss:.4f}, acc {acc:.4f}')

超参上,lr=1e-4配AdamW和weight_decay=1e-2是 Transformer 微调的稳妥起点,太大容易震荡,太小收敛慢。batch_size=32在 224 输入下对显存要求适中,显存不够就降到 16 并相应调小学习率。数据增强这里要特别小心:时频图的横轴是时间、纵轴是频率,水平翻转相当于把时间倒放,物理上不一定成立,所以我只给了 0.3 的低概率,垂直翻转(频率倒置)坚决不用。归一化用 0.5 均值方差是通用做法,如果你发现训练不稳,可以改成按数据集统计的真实均值和方差。

注意:预训练权重是在自然图像上训的,和时频图分布差异大,前几个 epoch 可以先用较小学习率 warmup,或者冻结 backbone 前几层,等 loss 稳定再解冻。

4. 把模型装进 GUI:从推理脚本到可交互界面

4.1 推理接口设计

GUI 和训练要解耦,推理部分单独封装成一个函数,输入是一段信号或一张图,输出是类别和置信度。这样 GUI 只负责调接口,不掺和模型细节。

import torch.nn.functional as F @torch.no_grad() def predict(model, img_tensor, classes, device='cuda'): """ img_tensor: 预处理后的张量, shape [1,3,H,W] 返回: (类别名, 置信度, 各类概率) """ model.eval() img_tensor = img_tensor.to(device) logits = model(img_tensor) probs = F.softmax(logits, dim=1)[0] conf, idx = probs.max(0) return classes[idx.item()], conf.item(), probs.cpu().numpy()

predict里用softmax把 logits 转成概率,方便 GUI 显示置信度条。classes从ImageFolder的class_to_idx反查得到,保证顺序一致。实际部署时把模型加载一次常驻内存,别每次点击都重新 load,否则界面会卡。

4.2 用 Tkinter 搭一个最小可用界面

import tkinter as tk from tkinter import filedialog, Label, Button from PIL import Image, ImageTk import torch from torchvision import transforms class BearingGUI: def __init__(self, model, classes, device='cpu'): self.model = model self.classes = classes self.device = device self.tf = transforms.Compose([ transforms.Grayscale(num_output_channels=3), transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.5]*3, std=[0.5]*3) ]) self.root = tk.Tk() self.root.title('轴承故障诊断 - WTFP + Swin') self.img_label = Label(self.root) self.img_label.pack() self.result_label = Label(self.root, text='请选择时频图', font=('Arial', 14)) self.result_label.pack(pady=10) Button(self.root, text='选择图片', command=self.load_image).pack(pady=5) def load_image(self): path = filedialog.askopenfilename(filetypes=[('PNG', '*.png'), ('JPG', '*.jpg')]) if not path: return img = Image.open(path).convert('L') disp = img.resize((224, 224)) self.tk_img = ImageTk.PhotoImage(disp) self.img_label.config(image=self.tk_img) tensor = self.tf(img).unsqueeze(0) cls, conf, _ = predict(self.model, tensor, self.classes, self.device) self.result_label.config(text=f'诊断结果: {cls} 置信度: {conf:.3f}') def run(self): self.root.mainloop() # 启动 # model = build_model(num_classes=10, device='cpu') # model.load_state_dict(torch.load('best_swin.pth', map_location='cpu')) # gui = BearingGUI(model, classes=['inner','outer','ball',...], device='cpu') # gui.run()

界面逻辑很直白:选图、显示、推理、出结果。Grayscale和训练时保持一致,否则输入分布对不上,结果会莫名其妙地差。device在 GUI 里建议用 CPU,避免和训练抢显存,单张推理 CPU 也就几十毫秒,体验够用。如果你想让界面更完整,可以加一个"批量诊断"按钮,遍历文件夹出 CSV 报告,这在现场巡检场景里很实用。

4.3 打包与部署的注意点

用 PyInstaller 打包时,timm和torch体积大,打出来的 exe 可能几百 MB,这是正常的。关键是模型权重文件要作为外部资源一起分发,别指望打进 exe 里还能方便替换。另外torchvision的版本要和torch匹配,否则加载模型时会报算子找不到。部署到工控机上,优先用 CPU 推理加 ONNX 导出,能省掉一大半环境依赖。

5. 避坑与排查:这套链路里最容易翻车的五个地方

5.1 时频图全黑或全白,模型学不到东西

现象是生成的 WTFP 要么一片黑要么一片白,训练 loss 不下降。原因通常是 dB 转换后没做归一化,或者信号幅值范围极端。解决方法是检查mag的数值范围,确保归一化前max - min不为零,必要时对信号先做去均值、除以标准差。另外log10里加了1e-12防止 log 零,这个不能省。

5.2 训练准确率很高但验证集崩盘

现象是训练集 99%,验证集 60% 上下。原因多半是数据泄漏——同一段信号切出来的重叠样本被分到了训练和验证两边。解决方法是按原始信号文件划分训练/验证,而不是按生成的图片随机分,保证同一段信号的所有切片只出现在一边。

5.3 类别不平衡导致模型只预测多数类

现象是混淆矩阵里少数类几乎全错。原因是 CWRU 里正常样本远多于故障样本。解决办法是用WeightedRandomSampler做重采样,或者在CrossEntropyLoss里传weight参数,按类别频率的倒数加权。

5.4 输入通道数不匹配报错

现象是RuntimeError: expected input to have 3 channels。原因是 WTFP 是单通道,而 Swin 预训练权重期望三通道。解决方法是 Dataset 里用Grayscale(num_output_channels=3)复制通道,别去改模型第一层,改了预训练权重就对不上了。

5.5 GUI 里结果和训练时对不上

现象是同一张图在脚本里预测正确,在 GUI 里预测错误。原因通常是预处理不一致,比如 GUI 里漏了归一化,或者classes顺序和训练时不同。解决办法是把预处理封装成同一个函数,训练和推理共用,classes从训练时保存的label_map.json读取,别手写。

6. 让这套方案真正好用:置信度校准与误报抑制的一个技巧

模型跑通只是第一步,现场最烦的是误报。Swin 输出的 softmax 置信度往往偏高,一个明显是噪声的样本也可能给出 0.9 的置信度,直接拿阈值卡会漏掉真正的早期故障。我一般会做两件事:一是温度缩放(temperature scaling),在验证集上拟合一个温度参数 T,把 logits 除以 T 再 softmax,让置信度更接近真实概率;二是对连续多段信号的预测做投票,单段误报被多段一致性过滤掉。

温度缩放的实现很简单,在验证集上最小化 NLL 即可:

import torch import torch.nn.functional as F import torch.optim as optim def calibrate_temperature(model, val_loader, device='cuda'): """在验证集上拟合温度参数 T""" model.eval() logits_list, labels_list = [], [] with torch.no_grad(): for imgs, labels in val_loader: imgs = imgs.to(device) logits_list.append(model(imgs).cpu()) labels_list.append(labels) logits = torch.cat(logits_list) labels = torch.cat(labels_list) T = torch.nn.Parameter(torch.ones(1) * 1.5) optimizer = optim.LBFGS([T], lr=0.01, max_iter=50) def closure(): optimizer.zero_grad() loss = F.cross_entropy(logits / T, labels) loss.backward() return loss optimizer.step(closure) return T.item() # 使用: 推理时 logits / T 再 softmax # T = calibrate_temperature(model, val_loader) # probs = F.softmax(logits / T, dim=1)

温度 T 一般落在 1.2 到 2.5 之间,T 越大输出越"谦虚"。校准后你会发现原本 0.95 的误报可能降到 0.6,配合一个 0.7 的阈值就能挡掉大部分。投票那块,我习惯对同一工况下连续 10 段信号分别预测,取众数作为最终结果,如果众数占比低于 60% 就标记为"疑似",交给人工复核。这套组合拳下来,现场误报率能压下去一大截。

最后说个我自己的习惯:每次换数据集或换工况,我一定先把 WTFP 图肉眼过一遍,确认冲击带清晰、类别之间有明显差异,再开始训模型。图都看不出区别,网络再深也是玄学。这套 WTFP + Swin 的链路我前后调了大半年,最大的教训就是别迷信模型结构,预处理和标签质量才是天花板。希望帮到你。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询