基于小波时频图与Swin Transformer的轴承故障诊断

简介：一份面向机械故障诊断领域研究者和工业设备运维工程师的完整项目实例文档，聚焦基于小波时频图与移位窗口视觉Transformer的轴承故障诊断方法。文档从原始振动信号切片、标准化出发，经连续小波变换构建时频图，再通过窗口注意力与移位窗口机制完成特征提取与故障分类，覆盖数据增强、模型训练、测试评估及部署应用全链路。整包为单个docx格式文档，大小147KB，内容含完整Python代码实现、GUI设计说明及代码逐段详解，便于对照学习与二次开发。目前已有63人学习，适合具备一定Python和深度学习基础，希望提升非平稳振动工况下诊断性能的读者参考。文档还重点解释了小波参数选择、类别不平衡及训练稳定性等工程挑战的解决方案，可帮助读者在风电、电机、轨道交通等场景快速落地智能运维系统。

1. 为什么一维信号不直接喂Transformer：小波时频图（WTFP）加移位窗口视觉Transformer才是轴承故障诊断的正解

做过轴承故障诊断的工程师大多有这种体验：采集卡里存了一整天的振动数据，FFT频谱翻来覆去看不出个所以然，手工提了十几个时域特征，内圈故障和外圈故障还是分不清。这套用Python实现的机械故障诊断方案换了条路：先把一维振动信号经连续小波变换生成小波时频图（WTFP），把非平稳的冲击特征变成图片上的亮斑和条纹，再用移位窗口视觉Transformer（Swin Transformer，简称ST）自动学图像特征做分类。整个方案一次跑通能拿到训练好的诊断模型和一个可交互的GUI工具，适合不想再做手工特征工程的算法工程师，也适合刚接触深度学习的机械专业学生照着复现。

很多人误以为Transformer只能处理一维序列或者自然语言，但Swin这类视觉Transformer在ImageNet分类上早就证明了自己；更关键的是，轴承故障信号里的冲击成分是典型的非平稳信号，天然适合先用小波变换把时间-频率联合分布画出来，再做图像分类。下面按“信号转图→图像分类→训练评估→GUI封装”的顺序，把这个项目的每个环节拆开讲。

2. 先把振动信号变成“故障照片”：小波时频图（WTFP）的生成原理与Python实现

2.1 为什么是时频图而不是直接喂原始波形：傅里叶的局限与Morlet小波的选择

轴承故障振动信号里的核心信息是周期性冲击：滚动体滚过缺陷时会产生一个衰减振荡，这个冲击在时间轴上只占几个毫秒，能量却集中在某个频带里。用FFT做全局频谱分析，相当于把这几个毫秒的瞬态冲击在整个采样时长里做平均，冲击特征被淹没在背景噪声里；这也是为什么很多人在频谱图上看到一堆峰值却找不到故障频率的原因。

短时傅里叶变换（STFT）能部分解决这个问题，它用固定长度的窗把信号切成片段再做FFT，但窗长一旦确定，低频和高频的时频分辨率就无法同时兼顾。连续小波变换（CWT）的思路完全不同：它用一组可伸缩、可平移的小波基函数和信号做内积，尺度小的高频小波窗口窄，能精确定位冲击发生时刻；尺度大的低频小波窗口宽，能捕捉趋势成分。这种“低频宽窗、高频窄窗”的变分辨率特性，天然匹配轴承冲击信号的特点。

具体选用哪种小波基，实践中90%的情况用Morlet小波就够了。Morlet是一个复值调制高斯波包，形态上像一个衰减振荡，和轴承故障的冲击响应波形高度相似；PyWavelets库里的小波名是'morl'，一行代码就可以调用。标题里写的WTFP，指的就是Wavelet Time-Frequency Plot，也就是把CWT系数矩阵的幅值渲染成二维灰度图，横轴时间、纵轴频率、像素亮度对应能量强度。这张图就是后面所有深度学习模型的输入。

2.2 用PyWavelets把一维信号转成时频图：核心代码与参数说明

以凯斯西储大学（CWRU）轴承数据集为例，这是轴承故障诊断领域最常用的公开数据集，驱动端振动信号采样率12kHz，故障包含内圈、外圈、滚动体和正常四类。先把原始mat文件读进来，截取固定长度，再用CWT变换生成时频图。

import numpy as np import pywt from scipy.io import loadmat def load_signal(mat_path, key='X096_DE_time', sample_len=4096): # 读取mat里的振动信号，取前sample_len个点作为一条样本 mat = loadmat(mat_path) data = mat[key].flatten() if len(data) > sample_len: start = np.random.randint(0, len(data) - sample_len) data = data[start:start + sample_len] return data def signal_to_wtfp(signal, fs=12000, num_scales=128): scales = np.arange(1, num_scales + 1) # 1~128尺度，对应频率从高到低 cwtmatr, freqs = pywt.cwt(signal, scales, 'morl', sampling_period=1.0/fs) power = np.abs(cwtmatr) # 复数系数取模 power_db = 20 * np.log10(power + 1e-12) # 转dB刻度压缩动态范围 return power_db, freqs

代码里有三个参数直接影响后续模型效果。第一个是sample_len，为什么取4096而不是1024或8192：12kHz采样率下轴承内圈故障的特征频率大概在100Hz到几百赫兹，4096个点能覆盖约0.34秒，包含几十个冲击周期，既保留足够模式又不会让单张图的计算量过大。第二个是num_scales，128个尺度决定了时频图的纵向像素数，太小会让频率分辨率不够，太大则图像尺寸膨胀、训练变慢，128是对精度和速度比较平衡的取值。第三个是转dB刻度这一步，CWT原始系数动态范围很大，最大值和最小值可能差几十个数量级，直接在matplotlib里保存会出现全黑或全白图，20*log10压缩是必须的，这里加了1e-12避免真零值取对数报错。

注意：MATLAB的cwt函数输出自带归一化，直接保存图片效果尚可；PyWavelets的cwt输出是裸系数，不做dB压缩和归一化就存图，后面模型几乎不可能收敛。这是换库迁移时最容易踩的坑。

2.3 图像尺寸、色彩映射与保存格式：影响模型精度的三个隐藏参数

生成时频图矩阵之后，把它落盘成图片这步看起来不起眼，却有三个参数学问很大。第一是目标尺寸，Swin Transformer的标准输入是224×224，直接把128行的power_db用双三次插值缩放到224×224即可；不建议用更大尺寸，模型参数量不变但计算量平方级上升，而故障时频图的判别纹理在缩放到224后依然清晰。

第二是色彩映射，matplotlib里写plt.imsave(path, power_db, cmap='jet')很顺手，但我的建议是保存成灰度图。jet这类彩色映射在视觉上确实更好看，但它本质是伪彩色编码，会把能量差的层次放大，而模型需要的是稳定的特征分布；灰度图用同样的0-255线性映射，规律更简单，也省去了推理时加载matplotlib的额外开销。真正影响的不是用哪种colormap，而是训练和推理时必须用同一种，保持一致这个动作比选型本身重要得多。

第三是文件格式，务必用PNG不要用JPG。JPG的有损压缩对自然照片是可行的，但对时频图这种充满细条纹和边缘的图，压缩噪声会被模型当成纹理学进去，实测中JPG质量因子85就能让验证精度掉两个百分点以上。下面是保存图片的代码，注意先把dB值归一化到0-255。

from PIL import Image def save_wtfp_as_image(power_db, save_path, target_size=(224, 224)): vmin, vmax = power_db.min(), power_db.max() img_8bit = ((power_db - vmin) / (vmax - vmin + 1e-6) * 255).astype(np.uint8) img = Image.fromarray(img_8bit, mode='L') # 灰度图 img = img.resize(target_size, Image.BICUBIC) img.save(save_path, format='PNG') return save_path

这段代码在保存前用min/max做了全局归一化。有个细节值得提：当样本之间存在明显能量差异时，全局min/max会让弱故障样本的对比度被强正常样本压低，所以更稳妥的做法是按5%和95%分位数截断后再归一化，牺牲一点极端值换取大多数样本的对比度稳定。图像保存好后，这个项目就从信号处理阶段进入了图像分类阶段。

3. 让Swin Transformer看得懂“故障照片”：移位窗口机制与数据准备

3.1 Swin Transformer相比ViT和CNN，为什么更适合时频图

普通视觉Transformer（ViT）把整张图切成16×16的patch，每个patch和全图所有patch计算全局注意力，在ImageNet大型数据集上效果不错，但直接用到时频图上有两个问题：一是计算复杂度随图像分辨率平方级增长，224×224的图切成小patch后序列长度很长，显存压力大；二是时频图的特点是局部密集纹理，一个冲击条带可能只覆盖几个patch，全局注意力会把注意力分散到不相关的区域，反而干扰局部模式提取。

Swin Transformer解决这两个问题的核心机制就是标题里说的“移位窗口”。它先把图像划分成7×7的小窗口，在每个窗口内部做自注意力，这一步控制了计算量；然后又引入shifted window操作，让下一次计算时窗口整体偏移一定像素，使得相邻窗口之间的信息能够互相流动。这种局部建模加跨窗口连接的结构，既保留了Transformer的建模能力，又接近CNN的局部感受野习惯，对时频图上“某个频带在某个时刻突然亮起”这种局部事件非常敏感。

还有一点是Swin自带的金字塔层级结构，每个阶段输出分辨率减半、通道翻倍，这和多尺度特征提取的思想不谋而合。轴承故障信息往往分布在多个频带，内圈故障特征频率及高次谐波都能在时频图上形成不同位置的亮斑，多尺度特征对这类“多频带联合判别”的任务有实打实的帮助。实际用的时候，选Swin-Tiny就够了，这个tiny版本参数约2800万，比ViT-Base小不少，在轴承故障这种类别少、模式相对固定的任务上精度和速度最平衡。

3.2 数据集目录组织与ImageFolder加载：最小可跑通的分割代码

训练一个图像分类模型，数据目录结构按torchvision的ImageFolder约定来组织是最省事的：根目录下每个类别一个文件夹，文件夹名就是类别名。对这个项目，建议按“正常/内圈故障/外圈故障/滚动体故障”四个目录组织训练数据。注意不要在目录名里带速度、载荷等工况信息，否则模型会把工况当作类别特征。

import os import random import shutil def split_dataset(src_root, dst_root, train_ratio=0.8, seed=42): # src_root下每个子目录是一类，按比例随机切分到train和val random.seed(seed) for cls_name in os.listdir(src_root): src_cls = os.path.join(src_root, cls_name) files = [f for f in os.listdir(src_cls) if f.endswith('.png')] random.shuffle(files) split_idx = int(len(files) * train_ratio) for split_name, split_files in [('train', files[:split_idx]), ('val', files[split_idx:])]: dst_cls = os.path.join(dst_root, split_name, cls_name) os.makedirs(dst_cls, exist_ok=True) for f in split_files: shutil.copy(os.path.join(src_cls, f), os.path.join(dst_cls, f))

这个脚本按类别分别洗牌再切分，保证train和val里四个类别的比例一致。如果你的原始数据里有多个工况（比如0马力、1马力、2马力负载），切分时一定要把同一段原始振动信号产生的所有窗口放进同一个集合，否则就出现数据泄漏，这一点到第5章避坑部分还会重点讲。切分完成后，用ImageFolder和DataLoader加载：

from torchvision import datasets, transforms from torch.utils.data import DataLoader transform_train = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(degrees=8), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) transform_val = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) train_ds = datasets.ImageFolder(root='data/train', transform=transform_train) val_ds = datasets.ImageFolder(root='data/val', transform=transform_val) train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4) val_loader = DataLoader(val_ds, batch_size=32, shuffle=False, num_workers=4)

ImageFolder会按文件夹名自动生成类别索引，训练时把索引到类名的映射保存下来，后面GUI推理要用。Normalize用的均值标准差是ImageNet的统计值，虽然时频图和自然图像的分布差别很大，但迁移学习时沿用这个设置是通行做法，不引入额外的调参负担。

3.3 数据增强的物理边界：旋转、翻转哪些能用哪些不能用

数据增强在自然图像任务里是可以大胆尝试的，但时频图有明确的物理坐标轴，横轴是时间、纵轴是频率，增强操作一旦破坏这个物理含义，模型就会学到错误的不变性。最典型的例子是垂直翻转，把时频图上下颠倒之后，原本1000Hz的频率变成了负频率方向上的1000Hz，这在物理上是不存在的信号，等于凭空造出一类假样本，模型看到这类样本只会学习“图片颠倒了”，推理时必翻车。

相比之下，水平翻转要安全得多：它相当于把振动信号在时间轴上反向播放，对稳定的轴承故障信号来说，冲击周期性在时间反转后依然存在，频域能量分布也不变，所以RandomHorizontalFlip可以用。随机旋转要控制幅度，±8到±10度问题不大，但旋转超过15度时，原本水平的等频率线明显倾斜，同样破坏了物理意义。还有一个经常会忽略的点：时频图的局部遮挡增强（Cutout或RandomErasing）是可以放心用的，因为遮挡掉某个时间段的能量，相当于模拟了信号中的瞬态丢失，这种增强方式物理上说得通，而且对防止过拟合相当有效。

我的实际配置是：Resize + RandomHorizontalFlip + RandomRotation(8) + ToTensor + Normalize，不做色彩抖动和随机缩放。色彩抖动会对灰度图做伪彩色扰动，反而引入不存在的频率能量差异；随机缩放改变的是频率刻度，也会干扰模型。一句话总结：给时频图做增强，先问自己“这个操作在物理世界里有对应含义吗”，没有就不能用。

4. 训练Swin Transformer：迁移学习、超参设置与评估指标

4.1 用ImageNet预训练权重做迁移学习：为什么“别人的照片知识”对振动图也有用

Swin Transformer的完整实现和预训练权重在timm库和官方开源仓库里都能拿到，用timm加载最方便。很多人质疑：ImageNet预训练权重学的是猫、狗、汽车的特征，拿来分类轴承故障时频图有意义吗？我的理解是，卷积网络或Transformer的浅层特征学的是边缘、角点、纹理基元这类通用模式，这些模式在时频图里同样存在——故障冲击在时频图上形成的就是不同方向的边缘和纹理条带。真正需要重新学的是高层语义组合方式，所以迁移学习的标准做法是替换分类头，只微调后面几层。

import timm import torch.nn as nn def build_swin_tiny(num_classes=4, pretrained=True): model = timm.create_model( 'swin_tiny_patch4_window7_224', pretrained=pretrained, num_classes=0 # 去掉原始分类头，只拿特征 ) in_features = model.num_features model.head = nn.Linear(in_features, num_classes) return model

这里把num_classes设成0让模型只输出特征向量，再自己接一个Linear分类头，这样能明确看到分类头的输入维度。Swin-Tiny的num_features是768，也就是说特征向量是768维，接四类分类器正好。加载预训练权重时如果遇到尺寸不匹配的报错，几乎都是分类头没去掉导致的，按上面的写法就不会有这个问题。

迁移学习还有一个策略选择：先冻结backbone只训练分类头，还是直接全量微调。常见做法是先冻结backbone训10个epoch，等分类头收敛后再解冻整个网络用更小的学习率微调。这个策略能避免一开始梯度太大把预训练特征冲坏，对训练数据量小于每类500张的情况尤其有效。

4.2 训练脚本：优化器、学习率、早停与混合精度的参数组合

训练Swin Transformer的配置和训练CNN有差异，最明显的差异在优化器。惯用的SGD+Momentum在Transformer上收敛偏慢，AdamW是标配，配合cosine学习率衰减和线性warmup。下面的代码是一个完整的训练循环骨架，包含混合精度和早停逻辑。

import torch import torch.nn as nn from torch.cuda.amp import autocast, GradScaler def train_model(model, train_loader, val_loader, epochs=30, device='cuda'): model.to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) scaler = GradScaler() best_acc = 0.0 patience, max_patience = 0, 10 for epoch in range(epochs): model.train() running_loss, correct, total = 0.0, 0, 0 for imgs, labels in train_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() running_loss += loss.item() correct += (outputs.argmax(1) == labels).sum().item() total += labels.size(0) train_acc = correct / total model.eval() val_correct, val_total = 0, 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels = imgs.to(device), labels.to(device) outputs = model(imgs) val_correct += (outputs.argmax(1) == labels).sum().item() val_total += labels.size(0) val_acc = val_correct / val_total if val_acc > best_acc: best_acc = val_acc torch.save(model.state_dict(), 'best_swin.pt') patience = 0 else: patience += 1 scheduler.step() print(f'epoch={epoch+1} loss={running_loss/len(train_loader):.4f} ' f'train_acc={train_acc:.4f} val_acc={val_acc:.4f}') if patience >= max_patience: print('early stop triggered') # 早停后回滚到best_swin.pt即可 break

这段代码里，学习率1e-4是微调预训练权重的安全起点；如果从零训练，学习率可以放到5e-4到1e-3。weight_decay设0.05是Swin官方配置的常用值，过大容易欠拟合，过小则Transformer训练后期不稳定。混合精度用torch的autocast和GradScaler组合，Swin-Tiny在225×224输入下即使batch_size=32，显存占用也在10GB上下，不开混合精度的话16GB显卡很紧张。早停的耐心值设10个epoch，轴承故障数据集类别少，模型通常20个epoch内就会收敛，耐心值设太大没有意义。下表是供参考的参数组合。

从零训练时学习率不能一上来就5e-4，一般前5个epoch做linear warmup从1e-5升到5e-4，否则前几步loss可能直接发散。timm的create_model已经内置了分类头但是随机初始化的，如果训练初期loss不降反升，先检查学习率，再把warmup加上。

4.3 混淆矩阵与逐类精度：比总准确率更重要的三个指标

训练完成后，总准确率只能告诉你“整体情况还行”，对于故障诊断项目这是远远不够的。设备运维的真实场景里，漏报一次内圈故障的代价可能远大于误报一次正常，这就必须看每一类的精确率、召回率和F1分数，还要用混淆矩阵定位哪些类别容易被搞混。

import numpy as np from sklearn.metrics import confusion_matrix, classification_report import seaborn as sns import matplotlib.pyplot as plt def evaluate_model(model, val_loader, class_names, device='cuda'): model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for imgs, labels in val_loader: imgs = imgs.to(device) outputs = model(imgs) all_preds.extend(outputs.argmax(1).cpu().numpy()) all_labels.extend(labels.numpy()) cm = confusion_matrix(all_labels, all_preds) print(classification_report(all_labels, all_preds, target_names=class_names)) plt.figure(figsize=(6, 5)) sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names) plt.xlabel('Predicted') plt.ylabel('Actual') plt.tight_layout() plt.savefig('confusion_matrix.png', dpi=150)

看混淆矩阵时重点看两个位置：一是主对角线，二是相邻类别之间的误报。轴承故障里最常见的是内圈故障和外圈故障互相混淆，因为两者的振动传递路径相似，时频图上冲击条带的位置接近；滚动体故障因为缺陷随保持架旋转，冲击幅值受载区调制，样本特征不稳定，召回率通常最低。如果内圈和外圈混淆严重，不要急着调模型，先回看时频图：采样点数是否太短、scales范围是否覆盖了故障特征频率所在的频带，这些信号预处理层面的问题比模型架构更致命。

5. 避坑排查：从信号预处理到模型收敛的5个典型翻车现场

5.1 时频图“糊成一片”，训练好几十个epoch准确率都在随机水平

现象：训练了30个epoch，训练准确率始终在25%上下徘徊，和随机猜测一模一样。打开保存的时频图，整张图要么几乎全黑只有一两条亮线，要么一片惨白看不到纹理。

原因：最常见的是CWT系数没做dB压缩就保存图片，原始系数动态范围太大，归一化时被少数几个极大值主导，多数能量细节被压缩到肉眼不可见的暗部；另一种情况是信号没做去均值，直流分量产生的零频能量把时频图的对比度彻底压垮。

解决：信号进入CWT之前先做标准化，减去均值除以标准差；CWT系数取绝对值后必须过20*log10压缩；归一化时不要用全局min/max，改用5%到95%分位数截断。修完这三处，重新生成图片，肉眼应该能看到清晰的频带条纹。

5.2 数据泄漏：验证集准确率99%，拿到现场数据立刻掉到70%

现象：训练过程中验证集准确率一路飙到99%，模型看起来完美；但把模型部署到新的设备数据上，准确率只有七成上下，工业现场根本不敢用。

原因：用滑窗切样本时窗口重叠度过高，同一段原始振动信号切出的多个窗口被分别分进了训练集和验证集。模型在训练时见过的“邻居窗口”出现在验证集里，等于开卷考试，验证成绩虚高。

解决：按“信号段”而不是“窗口”做数据集划分。先把连续采集的原始信号切成长段，比如每段10秒，再从每段内滑窗生成样本；划分训练集和验证集时保证整段信号只进一边，可以用sklearn的GroupShuffleSplit按段ID分组。验证集准确率会下降一些，但那个数字才是真实水平。

5.3 类别不平衡：正常样本是故障样本的十倍，模型把所有样本都判成正常

现象：正常类样本占90%，模型只学了一招“全预测成正常”就能拿到90%准确率，但对故障样本的召回率几乎是零，这种模型在现场等于摆设。

原因：交叉熵损失在多类不平衡时，大类的梯度完全压过小类，模型没有动力去学习少数类特征。

解决：第一个手段是给损失函数加权重，CrossEntropyLoss(weight=class_weights)，权重按各类别样本数的倒数归一化；第二个手段是数据过采样，对故障类样本在时频图层面做小幅旋转和水平翻转增强，把故障类扩到和正常类接近的比例。注意不要对故障类做大幅旋转，原因在第3章说过，会破坏物理合法性。

5.4 GPU显存不足：Swin-Tiny也能OOM，batch size直接起不来

现象：batch_size设32，跑了两个step就报CUDA out of memory，换Swin-Tiny还是爆显存。

原因：Swin-Tiny虽然参数只有2800万，但224×224输入的patch序列长度近3000，注意力机制的中间激活值显存占用远高于同等参数量的CNN。如果没开混合精度，8GB显卡跑batch_size=32几乎没有可能。

解决：先开autocast混合精度，显存直接减半；还不行就把batch_size降到16或8，用梯度累积模拟大batch，accum_steps设为4时每4个小batch更新一次权重，效果等效于batch_size=32。最后的手段是把输入分辨率从224降到192，Swin-Tiny的窗口设计不依赖特定分辨率，降分辨率对精度影响通常不超过1个百分点。

5.5 迁移学习效果反而比从零训练差：预训练权重不是万能金药

现象：用了ImageNet预训练权重微调，验证集准确率85%；同事同样数据从零训练，反而到92%。预训练权重成了负优化。

原因：ImageNet特征来自自然图像的空间纹理，而时频图的纹理统计特性差异很大，比如频率轴方向的条纹有很强的方向性，这是自然图像里少见的结构。当训练数据量足够多时，从零训练可以完全依据时频图的数据分布学特征，反而更贴合目标域。

解决：数据量每类超过1000张时，把“是否用预训练”当成一个超参，两种都试一遍再选验证集高的，不要默认迁移学习一定好。每类只有几百张时，预训练还是值得用的，但需要保证warmup足够充分。这个边界本身有点玄学，和你选择的信号长度、scales范围都有关系，唯一的确定性做法是跑两组对比实验。

6. GUI设计阶段：把模型封装成工程师愿意打开的诊断工具

6.1 界面选型：为什么我选Tkinter而不是PyQt

训练出的模型再好，如果交付给现场工程师的是一个Jupyter Notebook，基本等于没做。诊断工具的界面设计不必花哨，但要足够直观。Tkinter是Python标准库自带，打包后的exe体积小，做“选择文件→显示时频图→输出故障类型”这种单页面工具足够用；PyQt5界面更现代，但打包体积和依赖复杂度都上一个量级。我个人倾向Tkinter，因为它跨平台稳定、不依赖额外授权，运维人员双击就能用。

6.2 GUI推理链路：从按钮回调到后台线程，核心代码与模型加载

GUI的核心逻辑只有三步：读取振动信号、生成时频图、模型预测。但有个工程坑必须处理：时频图生成和模型推理都比较耗时，如果把这两步直接放在按钮回调里，界面会卡死好几秒，用户会以为程序崩溃。正确做法是把推理放进独立线程，主线程保持界面响应。

import threading import tkinter as tk from tkinter import filedialog, Label, Button import torch from PIL import Image import numpy as np import pywt class DiagnosticGUI: def __init__(self): self.root = tk.Tk() self.root.title('轴承故障诊断 - Swin Transformer') self.model = self.load_model('best_swin.pt') self.result_label = Label(self.root, text='请选择振动数据文件', font=('Arial', 14)) self.result_label.pack(pady=10) Button(self.root, text='载入数据并诊断', command=self.start_predict, width=20).pack(pady=5) def load_model(self, weight_path): import timm import torch.nn as nn model = timm.create_model('swin_tiny_patch4_window7_224', pretrained=False) model.head = nn.Linear(model.num_features, 4) model.load_state_dict(torch.load(weight_path, map_location='cpu')) model.eval() return model def start_predict(self): path = filedialog.askopenfilename(filetypes=[('mat files', '*.mat')]) if path: threading.Thread(target=self.predict, args=(path,), daemon=True).start() def predict(self, mat_path): signal = load_signal(mat_path) # 复用第2章的加载函数 power_db, _ = signal_to_wtfp(signal) # 复用第2章的变换函数 img = save_and_load_image(power_db) # 归一化并缩放为224x224 img_tensor = torch.from_numpy(np.array(img)).float() / 255.0 img_tensor = img_tensor.unsqueeze(0).repeat(1, 3, 1, 1) # 灰度图复制成3通道 img_tensor = (img_tensor - torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)) / \ torch.tensor([0.229, 0.224, 0.225]).view(3,1,1) with torch.no_grad(): pred = self.model(img_tensor.unsqueeze(0)) probs = torch.softmax(pred, dim=1).squeeze(0).numpy() cls_names = ['正常', '内圈故障', '外圈故障', '滚动体故障'] idx = int(probs.argmax()) msg = f'诊断结果: {cls_names[idx]} 置信度: {probs[idx]*100:.2f}%' self.result_label.config(text=msg)

加载模型时用pretrained=False并直接load_state_dict，避免推理机还需要ImageNet预训练权重。灰度图复制成三通道再按ImageNet统计量归一化，这一步和训练时的预处理必须完全一致，否则模型精度会莫名其妙下降。GUI代码里没有定义load_signal和signal_to_wtfp，实际项目把它们放在独立模块里复用，避免复制粘贴导致预处理逻辑漂移。

我最初做这个工具时把所有逻辑都塞进按钮回调，点一次卡五秒，被现场工程师吐槽是“老年机软件”。后来自学线程改造，界面流畅了，但当时因为训练和推理预处理不一致，GUI里出的结果和离线测试差了一大截，排查了一整天才发现是归一化顺序错了。这个血泪经验提醒我：训练脚本和GUI推理必须共享同一个信号处理模块，不要各写各的；模型上线前拿几条已知故障样本走一遍GUI全流程做回归验证，比什么部署技巧都管用。希望帮到你。

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