简介这份资源面向具备Python与深度学习基础的科研人员、研究生及工业设备状态监测工程师提供一套基于小波时频图WTFP结合移位窗口视觉TransformerST的轴承故障诊断完整项目实例。它针对非平稳振动、工况变化与跨设备迁移下传统方法诊断性能不佳的问题构建了从原始信号切片、标准化、连续小波变换生成时频图、数据增强、模型训练到测试评估的端到端链路并覆盖GPU推理、服务接口与监控维护等部署环节。资源包为1个docx文档约147KB内含完整程序代码、GUI设计与逐段代码详解目录涵盖项目背景、模型架构、窗口注意力与移位窗口模块、训练验证及测试评估等章节。已有63人学习。读者可据此复现可运行的故障识别原型掌握小波参数选择、数据划分与训练稳定性等工程要点为风力发电、工业电机、数控机床等旋转机械的智能运维提供可落地的技术参考。1. 轴承故障诊断为什么值得用「小波时频图 移位窗口 Transformer」重做一遍振动信号是旋转机械最诚实的体检报告但这份报告是时变的、非平稳的轴承一旦出现早期点蚀或剥落冲击成分往往埋在强噪声和工频干扰里。传统做法要么靠人工特征峭度、包络谱、均方根加浅层分类器要么直接把一维时序丢进 CNN结果在变工况、变转速下泛化能力很差。这几年我反复试过一条路线先把一维振动信号转成小波时频图WTFP把时间-频率的二维纹理显式暴露出来再用移位窗口视觉 TransformerSTSwin Transformer 那一类结构去做分类。它解决的核心问题是——让模型既能看到局部冲击的细节又能建模全局的频带依赖同时把计算量压到工程可接受的范围。这套方案适合做设备状态监测、故障诊断方向的工程师和研究生尤其是手里已经有 CWRU 这类公开轴承数据集、想快速跑通一条完整链路的人。下面我按信号怎么变图、图怎么进网络、网络怎么落地成 GUI的顺序把能抄作业的部分全写出来。2. 从一维振动到小波时频图WTFP 的生成链路与参数选择2.1 为什么选小波而不是 STFT短时傅里叶变换STFT的窗长一旦定死时间分辨率和频率分辨率就固定了而轴承故障的冲击成分频率高、持续时间短低频的转频调制又需要长窗。小波变换用可伸缩的母小波高频处窗窄、低频处窗宽天然适配这种多尺度需求。工程上我一般用连续小波变换CWT生成时频图母小波优先选 Morlet 或复 Morletcmor因为它在时频聚集性上表现稳定做出来的图纹理清晰喂给视觉模型时特征更好认。选型上还有几个现实考量一是 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, fs12000, scalesNone, waveletcmor1.5-1.0, img_size(224, 224), save_pathNone): 将一维振动信号转成小波时频图 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_period1/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, fs12000, save_pathwtfp_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_len2048, stride1024): 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(parentsTrue, exist_okTrue) 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:startseg_len] # signal_to_wtfp(seg, save_pathsave_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_asciiFalse, indent2) 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_classes10, model_nameswin_tiny_patch4_window7_224, pretrainedTrue, drop_rate0.2): super().__init__() # timm 直接给出现成 Swin, 改分类头即可 self.backbone timm.create_model( model_name, pretrainedpretrained, num_classesnum_classes, drop_ratedrop_rate ) def forward(self, x): return self.backbone(x) def build_model(num_classes10, devicecuda): model BearingSwin(num_classesnum_classes) model.to(device) return model # 快速验证输入输出 # model build_model(num_classes10, devicecpu) # 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_size32, img_size224): tf transforms.Compose([ transforms.Grayscale(num_output_channels3), # 单通道转三通道 transforms.Resize((img_size, img_size)), transforms.RandomHorizontalFlip(p0.3), # 时频图翻转要谨慎 transforms.ToTensor(), transforms.Normalize(mean[0.5]*3, std[0.5]*3) ]) ds datasets.ImageFolder(data_dir, transformtf) loader DataLoader(ds, batch_sizebatch_size, shuffleTrue, num_workers4, pin_memoryTrue) return loader, ds.classes # 训练主流程 # device cuda if torch.cuda.is_available() else cpu # model build_model(num_classes10, devicedevice) # loader, classes build_loaders(wtfp_dataset) # criterion nn.CrossEntropyLoss() # optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-2) # for epoch in range(50): # loss, acc train_one_epoch(model, loader, criterion, optimizer, device) # print(fepoch {epoch}, loss {loss:.4f}, acc {acc:.4f})超参上lr1e-4配AdamW和weight_decay1e-2是 Transformer 微调的稳妥起点太大容易震荡太小收敛慢。batch_size32在 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, devicecuda): img_tensor: 预处理后的张量, shape [1,3,H,W] 返回: (类别名, 置信度, 各类概率) model.eval() img_tensor img_tensor.to(device) logits model(img_tensor) probs F.softmax(logits, dim1)[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, devicecpu): self.model model self.classes classes self.device device self.tf transforms.Compose([ transforms.Grayscale(num_output_channels3), 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(pady10) Button(self.root, text选择图片, commandself.load_image).pack(pady5) 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(imageself.tk_img) tensor self.tf(img).unsqueeze(0) cls, conf, _ predict(self.model, tensor, self.classes, self.device) self.result_label.config(textf诊断结果: {cls} 置信度: {conf:.3f}) def run(self): self.root.mainloop() # 启动 # model build_model(num_classes10, devicecpu) # model.load_state_dict(torch.load(best_swin.pth, map_locationcpu)) # gui BearingGUI(model, classes[inner,outer,ball,...], devicecpu) # 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_channels3)复制通道别去改模型第一层改了预训练权重就对不上了。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, devicecuda): 在验证集上拟合温度参数 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], lr0.01, max_iter50) 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, dim1)温度 T 一般落在 1.2 到 2.5 之间T 越大输出越谦虚。校准后你会发现原本 0.95 的误报可能降到 0.6配合一个 0.7 的阈值就能挡掉大部分。投票那块我习惯对同一工况下连续 10 段信号分别预测取众数作为最终结果如果众数占比低于 60% 就标记为疑似交给人工复核。这套组合拳下来现场误报率能压下去一大截。最后说个我自己的习惯每次换数据集或换工况我一定先把 WTFP 图肉眼过一遍确认冲击带清晰、类别之间有明显差异再开始训模型。图都看不出区别网络再深也是玄学。这套 WTFP Swin 的链路我前后调了大半年最大的教训就是别迷信模型结构预处理和标签质量才是天花板。希望帮到你。本文还有配套的精品资源点击获取