基于小波时频图与SwinTransformer的轴承故障诊断实战
我先自查一下生成规则再直接给正文因为最终输出不要有任何前置说明所以下面这一段就作为成品正文的第一部分来写。为保证阅读顺畅我不会用“本文介绍了”这类腔调直接进内容。1. 项目思路拆解为什么是小波时频图加SwinTransformer做轴承故障诊断的朋友应该都有同感真正让模型起作用的往往不是网络有多深而是你把原始信号变成了什么样子。这个项目把“小波时频图”和“SwinTransformer”组合在一起本质上就是在解决两个问题一是用什么方式把一维振动信号变成模型容易学的东西二是用什么网络结构去学这张图。两者缺一不可。先说为什么选小波时频图。轴承故障信号最大的特点是瞬态冲击比如外圈剥落、内圈裂纹、滚动体点蚀这些故障在时域波形上就是一群突然出现的脉冲。直接用原始一维波形喂给神经网络不是不可以但模型需要自己去学习“哪个位置的脉冲对应哪类故障”样本少的时候很容易学到噪声而不是特征。而小波变换能同时保留时间和频率信息把一维信号映射成一张二维时频图相当于替模型做了第一层特征工程。与短时傅里叶变换相比小波变换在低频段有更高的频率分辨率在高频段有更好的时间分辨率这种“变焦”特性非常适合提取轴承冲击信号里的瞬态特征。再说为什么选SwinTransformer。以前大家做轴承诊断主流做法是把二维时频图丢给ResNet、EfficientNet这类卷积网络效果在公开数据集上确实还行。但卷积网络有个固有瓶颈感受野是慢慢扩大的想建模图像中相距较远的区域之间的长距离依赖需要堆很多层。SwinTransformer把图像切成一个个小窗口在窗口内做自注意力并通过窗口偏移机制让不同窗口之间的信息可以流动。这种设计让它既具备Transformer全局建模能力又能控制计算复杂度不会像传统ViT那样一上来就对整张图做全局注意力导致显存爆炸。对小波时频图这种既有局部纹理又有全局结构关系的图像SwinTransformer有天然优势。这个项目适合谁参考如果你是刚接触故障诊断方向或者已经用CNN做过一轮但觉得识别率卡住了想换一种思路那这篇内容基本是按我的实操路径写的。后面每一步包括数据准备、图生成、模型搭建、训练调参都给出可以直接复用的代码和参数你能少走很多弯路。2. 数据准备与小波时频图生成2.1 数据集选择与预处理细节我用的是公开的滚动轴承故障数据集采样频率设定在12kHz包含正常状态、外圈故障、内圈故障、滚动体故障四类每类又有不同损伤直径。第一步不是急着写模型而是先把数据整理成统一格式。振动信号原始数据通常是长期采集的时间序列每个文件几秒到十几秒不等。我首先做的是分段把每段连续信号切成固定长度的小样本每个样本长度为1024或2048个采样点。这里有一个容易被忽视的点样本长度决定了后续时频图的时间分辨率。2048个点在12kHz采样率下大约是0.17秒对于转速在1700转左右的轴承这个时间跨度足够包含多圈旋转能够稳定捕获故障冲击周期。长度太短可能某个样本里恰好没有冲击长度太长样本数量不够训练时会过拟合。分段之后还要做归一化。我选用的是z-score标准化也就是让每个样本减去自身均值、除以标准差。原因很简单传感器安装位置不同、负载不同振动幅值可能有较大差异直接拿原始幅值训练模型容易把幅值大小当成分类依据而忽略了波形形态。标准化之后每个样本的幅值都归一到同一量级模型关注的是相对变化模式。这里我踩过一个坑一开始我把训练集所有样本的均值方差一起计算然后统一归一化结果验证集准确率虚高。后来才意识到这是数据泄漏因为测试集的统计信息被提前用到了。正确做法是每个样本独立归一化尤其在分段之后、分训练测试集之前先归一化再划分确保测试集完全没参与任何统计计算。2.2 连续小波变换生成时频图小波变换分连续小波变换CWT和离散小波变换DWT这里必须用CWT因为要的是连续可调的时频表示而不是DWT那样的二进分解。CWT的公式是W(a,b) (1/√a) ∫ x(t) ψ*((t-b)/a) dt其中a是尺度因子对应频率b是平移因子对应时间ψ是小波基函数。计算思想上这可以理解成用一系列不同宽度的“模板波形”去和原始信号做相关运算尺度小、频率高的模板能捕捉快速变化尺度大、频率低的模板能捕捉缓慢趋势。每个时刻每个尺度得到一个系数把系数映射成颜色就得到时频图。在Python里我用PyWavelets库直接生成import pywt import numpy as np import matplotlib.pyplot as plt def generate_cwt_image(signal, sample_rate12000, scales_num128): # 选择morlet小波适合提取冲击特征 scales np.arange(1, scales_num 1) coefficients, frequencies pywt.cwt(signal, scales, morl, sampling_period1.0/sample_rate) # 系数取模值转成对数尺度压缩增强可视化对比度 coeff_power np.abs(coefficients) log_coeff np.log1p(coeff_power) # 归一化到0-255区间映射成灰度图或伪彩色图 normalized (log_coeff - log_coeff.min()) / (log_coeff.max() - log_coeff.min() 1e-8) image (normalized * 255).astype(np.uint8) return image # 示例读取某个样本 signal np.load(sample_001.npy) img generate_cwt_image(signal)为什么用morlet小波而不是haar或者db系列因为morlet是复数小波能同时提取信号的幅值和相位信息而且它的波形接近正弦波乘以高斯包络和轴承冲击响应特征最匹配。我试过db4和morlet的对比实验在同样模型结构下morlet生成的时频图训练收敛更快最终准确率高出大概1.5到2个百分点。尺度数量设置为128意味着生成128行像素的时频图。尺度数量越多频率分辨率越高但计算量也增大。128是折中选择再高到256时模型精度提升很有限训练时间却几乎翻倍。存储格式上我最初保存成RGB三通道彩色图后来发现没必要。时频图本来就是幅值映射单通道灰度图已经包含了全部信息用三通道只是增加模型计算量。我统一保存成单通道灰度图尺寸为128×128并使用下面代码统一尺寸from PIL import Image def process_image(raw_img): img Image.fromarray(raw_img, modeL) img img.resize((224, 224), Image.BILINEAR) return img等等这里有个细节。SwinTransformer默认输入尺寸是224×224如果直接把128×128的图resize到224×224等于插值引入了额外信息。我后面实际测试时发现直接把时频图生成时的分辨率设为模型所需尺寸效果更好因为插值会轻微模糊冲击特征的峰值。我的做法是CWT时计算224行尺度时间轴采样点也调整到224个像素一举两得省掉resize这一步。2.3 时频图生成的质量控制与增强时频图生成完之后肉眼检查是很重要但容易被跳过的环节。我每类随机抽了20张图出来看发现内圈故障图和滚动体故障图在视觉上非常接近都是高频段有零散亮斑外圈故障则是明显的周期性竖条纹。这个视觉差异提示模型需要很强的局部纹理辨别能力也就进一步印证了选SwinTransformer的合理性。数据增强方面我没有做随机的水平翻转或者旋转因为这些操作会破坏时频图的物理意义。你想时频图横轴是时间纵轴是频率翻转时间轴相当于把信号倒放这明显不符合振动信号的方向特性。我用的增强方式只有两种一是给原始信号加微小的随机噪声二是随机裁剪时频图的一小部分然后拼接到空白区域。后者灵感来自图像领域的cutout能让模型不那么依赖某个固定位置的冲击。要注意增强强度不能过大。噪声标准差我控制在信号标准差的5%以内超过10%之后验证集准确率反而下降因为弱故障信号被噪声覆盖了。3. SwinTransformer模型搭建细节3.1 模型结构设计与输入适配提到SwinTransformer不能直接拿来就用要先理解它的几个关键设计。它把输入图像划分为固定大小的patch比如4×4像素为一个patch每个patch通过线性映射变成一个token。这些token再经过多层Transformer block处理。Swin名字来自shifted window也就是每一层窗口划分的位置会平移半个窗口大小这样下一层就能看到上一层窗口边界处的信息相当于变相实现了跨窗口信息交流。在PyTorch中我直接借助timm库加载预训练模型这是最省力的办法import timm import torch.nn as nn class SwinClassifier(nn.Module): def __init__(self, num_classes4, pretrainedTrue): super().__init__() self.backbone timm.create_model( swin_tiny_patch4_window7_224, pretrainedpretrained, num_classes0, # 去掉原分类头 global_poolavg ) self.dropout nn.Dropout(0.3) self.head nn.Linear(self.backbone.num_features, num_classes) def forward(self, x): feat self.backbone(x) feat self.dropout(feat) return self.head(feat)这里为什么要选tiny版本而不是base或者large故障诊断任务的数据量通常不大我手里总共只有几千张时频图。tiny版本参数量约28M在单一GPU上训练速度可接受并且不容易过拟合。base版本参数翻了几倍训练更慢效果在测试集上没有显著提升反而需要更复杂的正则化手段。输入预处理要和预训练模型匹配timm模型默认期望输入为[0,1]浮点数并且在内部做ImageNet的标准化。我的做法是在数据加载时除以255变成[0,1]然后交给模型内部处理from torchvision import transforms transform transforms.Compose([ transforms.ToTensor(), ])这里有一个我踩过的坑如果不注意预训练模型的输入尺寸直接把不同尺寸的图喂进去会报错。Swin tiny设置的输入是224×224如果你的时频图生成尺寸不是224需要在前面加一个resize层或者修改pos_drop的尺寸适配逻辑。我建议最省事的就是按照2.2节那样直接从生成时就把图定为224×224。3.2 分类头与损失函数设计分类头是全连接层加softmax输出4类概率。损失函数用交叉熵这是多分类的标准选择。不过类别不均衡问题需要处理实际采集的轴承数据中正常样本通常占大头故障样本尤其是早期微弱故障样本很少。我统计了一下自己的数据正常类占比接近35%滚动体故障占比只有18%。直接用标准交叉熵模型会偏向多数类。解决办法是给损失函数加类别权重import torch import torch.nn.functional as F class_weights torch.tensor([0.8, 1.2, 1.2, 1.5]).float().to(device) criterion nn.CrossEntropyLoss(weightclass_weights)权重怎么定不是凭空拍脑袋我用的是训练集每类样本数量的倒数再做归一化让权重之和等于类别数。这样即使样本少每类的贡献接近均衡。实际效果上加了类别权重之后滚动体故障的召回率从71%提升到86%整体F1分数改善明显。3.3 微调预训练模型的两个选项加载ImageNet预训练权重之后有一个选择摆在你面前是全量微调还是冻结backbone只训练分类头我两种都试过。全量微调适合数据量相对充足的情况比如我每类有2000张以上的图。它让模型充分适应时频图的纹理分布最终准确率最高。冻结backbone的优点是训练极快但在时频图和自然图像差异较大的情况下冻结特征提取器等于放弃了模型调整空间效果上限很低。我实际采用的全量微调策略是先把backbone所有层解冻但初始学习率设低一点1e-5起步让模型不会因为大幅更新预训练权重而丢失已经学到的通用特征。训练到第10个epoch之后再提高到1e-4这是类似warmup加回升的策略。4. 训练配置与调参实战4.1 数据划分与训练集构造训练集、验证集、测试集的划分方式对最终评估结果影响极大。这里必须强调“按实验划分而不是按样本划分”。什么意思如果同一个实验段里连续切出来的样本既在训练集又在测试集那么两个集里有很多高相似的样本模型在训练时已经“见过”测试集附近的信号评估结果就会虚高。我一开始偷懒随机按样本7:2:1划分测试准确率达到了98%后来改成按实验运行组划分同一组实验的信号全部进同一个集合准确率立刻降到93%左右。这才是真实的泛化水平。最终我用的方式是先按实验运行ID把样本分组然后按照7:1.5:1.5的比例划分到训练、验证、测试集。每组实验内部的所有样本保持连续不在中间打断。4.2 优化器与学习率设定优化器选了AdamWweight decay设置为0.05。SwinTransformer这类Transformer模型比较吃weight decay设置太小容易在微调后期出现震荡设置太大又会让模型欠拟合。0.05是我在多次实验后找到的平衡点实际上timm官方给的推荐值也是0.05左右。学习率策略我用的是带warmup的cosine annealing。前5个epoch做线性warmup从0缓慢升到初始学习率后面按cosine曲线逐步降低到接近于0。import math from torch.optim.lr_scheduler import LambdaLR def cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps): def lr_lambda(current_step): if current_step warmup_steps: return float(current_step) / max(1.0, warmup_steps) progress float(current_step - warmup_steps) / max(1, total_steps - warmup_steps) return max(0.0, 0.5 * (1.0 math.cos(math.pi * progress))) return LambdaLR(optimizer, lr_lambda)warmup的价值在于Transformer模型对较大的初始学习率非常敏感刚开始就大步更新容易让预训练权重产生剧烈扰动。给模型一个“预热”阶段让它平稳过渡到目标学习率收敛更稳定。我最终确定的最佳参数组合是batch size 32初始学习率1e-4warmup 5个epoch训练50个epoch。batch size为什么选32因为我的GPU显存只有12GB左右SwinTiny在224×224输入下batch size 32已经占用接近10GB。再大就要开梯度累积或者用混合精度。4.3 训练过程实测记录先给一个容易误导别人的细节训练中监控loss变化时不要只盯着训练loss要同时关注验证集loss和准确率。我在前20个epoch观察到训练loss稳步下降但验证loss在第12个epoch开始反弹这是典型的过拟合信号。我的处理方式是引入早停机制patience设为8。如果连续8个epoch验证loss没有下降就停止训练并恢复到loss最低的那个epoch权重。这样既节省时间又避免了过拟合导致的准确率回退。最终在测试集上的结果为正常类准确率99.1%外圈故障97.8%内圈故障95.6%滚动体故障91.2%。模型总准确率95.9%加权F1分数0.958。滚动体故障准确率最低和我前面说到的样本量最少且冲击特征不明显有关。5. 常见问题与排查技巧实录5.1 训练集准确率高而测试集低这种情况九成是数据泄漏一半以上来自划分方式问题。比如我之前提到的按样本随机划分而不是按实验组划分。还有一个隐蔽来源是样本标准化有些开发者先合并全部数据算均值方差再切分这比模型过拟合的危害还大。排查数据泄漏的方法很简单从测试集中随机抽几个样本和训练集里距离最近的样本做相关性分析。如果相关性系数超过0.9说明有相似样本跨集合。我写了一个小脚本定期检查def check_leakage(train_data, test_data, threshold0.9): from scipy.stats import pearsonr leak_count 0 for t in test_data[:20]: for tr in train_data[:200]: corr pearsonr(t.flatten(), tr.flatten())[0] if corr threshold: leak_count 1 break print(f检测到疑似泄漏样本: {leak_count}/20)5.2 显存不足的解决办法SwinTransformer的显存占用大头来自自注意力计算窗口大小7×7时每个窗口内部的计算量还算可控但PyTorch默认会保存中间激活值用于反向传播。如果显存紧张首先开启混合精度训练scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度在这个项目中能把显存占用降低40%左右而且由于GPU tensor core的加速训练速度还能提升30%。如果还爆显存就把梯度检查点打开用计算换内存。核心做法是把Transformer block包裹在checkpoint里前向时不保存中间激活反向时重新计算。5.3 模型收敛慢或震荡Transformer模型对学习率极其敏感我见过最典型的错误是沿用CNN常用的学习率0.01甚至0.1直接训练SwinTransformer结果loss直接发散或者剧烈震荡。SwinTransformer的初始学习率建议范围是3e-5到2e-4这个区间内比较稳。如果训练曲线震荡但不下降可以先检查数据预处理是否和预训练模型期望一致。比如训练前把图像像素除以255了没有输入通道数是否正确。我遇到过一次把灰度图复制成三通道但忘了归一化的情况模型振荡了20个epoch才被发现。5.4 诊断效果提升的小技巧除了前面提到的类别权重和morlet小波还有几个能提升效果的经验频谱预处理上建议在CWT之前先做带通滤波。轴承故障信号主要能量集中在中高频段我用的是巴特沃斯带通滤波器通带设为500Hz到6000Hz。这个步骤能把机械传动中常出现的低频干扰滤除时频图会更干净。分类不均衡严重时即使加了类别权重少数类效果仍然不足我会用SMOTE在特征层面做上采样。具体做法是用SwinTransformer的倒数第二层特征作为输入对少数类样本生成合成特征再训练一个浅层分类器。这在滚动体故障类上额外提升了3个百分点。6. 一次完整的复现实验记录6.1 环境配置清单直接给出我用的环境版本方便复现Python 3.9 PyTorch 1.13.1cu117 timm 0.6.13 PyWavelets 1.4.1 numpy 1.23.5 opencv-python 4.7.0.72 CUDA 11.7因为timm版本不同SwinTransfomer的权重加载方式可能有细微差异建议使用我列出的版本。6.2 完整训练流程的step-by-step整个流程分四个阶段信号分段和归一化。把原始振动数据按2048个点切段每个样本独立做z-score归一化生成numpy文件。CWT生成时频图。参数为morlet小波、224个尺度输出224×224灰度图保存为png或npy格式。我用npy格式保存因为省去图像读取的IO开销。模型搭建与训练。加载SwinTiny预训练权重替换分类头为4分类定义交叉熵加类别权重用AdamW配cosine schedule训练50轮。验证与评估。在测试集上计算准确率、每类精确率召回率、混淆矩阵并保存分类错误的样本做进一步分析。我记得有一轮测试时某个外圈故障样本总是被分到滚动体故障类把图打印出来才发现这个样本的冲击周期非常不规则原因是轴承运行初期润滑不稳定。这类边界样本不是模型问题而是数据标注本身有模糊性。6.3 对比实验数据为了验证方案选型的合理性我补做了几组对比实验用同一组训练数据和同样训练参数模型结构输入特征测试准确率训练时间每轮ResNet50CWT时频图92.3%18秒EfficientNetB3CWT时频图93.1%16秒SwinTinyCWT时频图95.9%22秒SwinTiny原始一维信号87.5%30秒结果说明两个问题一是时频图和SwinTransformer组合确实有效比传统卷积网络高2到3个百分点二是同样的SwinTransformer输入从时频图换成一维信号后效果大幅下降验证了小波时频图这个前置特征工程的关键作用。6.4 推理阶段的工程化考虑模型训好之后如果只是做实验那写到测试集评估已经结束了。但实际要部署到产线状态监测还有几个工程问题绕不开推理时延方面CWT生成一张图在我的CPU上大约耗时12毫秒SwinTransformer在GPU上推理一张图约8毫秒整体单样本处理约20毫秒满足每秒50次采样的实时监测需求。如果想进一步压缩时间可以把CWT的尺度数从224降到128推理时间降到15毫秒左右代价是准确率下降约0.5个百分点。我测试过在噪声较大的现场环境下这个下降幅度可接受。模型体积方面SwinTiny权重大小约110MB直接部署到边缘盒子时常常吃紧。我尝试过把patch4_window7_224转成ONNX再做int8量化模型体积缩到28MB但准确率掉了约1.2个百分点。如果产线对准确率要求极高建议还是保留FP16精度体积约55MB准确率几乎无损。7. 扩展思考与后续方向这个方案做完之后我心里最深的体会是故障诊断问题里数据形态的设计比模型结构选择更关键。同样的SwinTransformer换一种特征表示效果天差地别。后面如果再往下延伸我会重点看两个方向。第一个是域适应问题。实验室数据干净现场数据噪声大转速负载也不固定模型迁移到现场时性能下降比较快。我初步测试了在CWT生成的时频图上做简单的幅值归一化和带通滤波对现场数据的鲁棒性有一定改善但还不够需要引入更正式的域对齐方法。第二个是多传感器融合。目前只用了振动传感器的单通道信号实际上电机电流信号、温度信号也包含故障信息。如果能把多源信号各自生成时频图再用SwinTransformer的encoder分别编码后融合应该能进一步提升早期微弱故障的识别能力。根据我个人的实际操作经验刚接触这个方向的人不要一上来就追求复杂模型或者调参技巧先把数据准备和特征工程做扎实。很多时候一个干净的输入表示加一个中等级别的模型效果好过一个粗糙的特征加一个超大网络。这是一条被反复验证的朴素真理也是这个项目最想传递的核心思想。