在旋转机械故障诊断这个方向上轴承的振动信号分析永远是绕不开的话题。我之前做过几轮基于传统特征提取的项目也试过直接喂一维信号给卷积分网络最终在实际项目里稳定复现、效果也最让我放心的方案是小波时频图加SwinTransformer这套组合。简单说就是先把一维振动信号做连续小波变换变成二维时频图再交给SwinTransformer做分类。整个流程基于Python和PyTorch实现代码结构清晰训练效率和识别精度都不错。这个思路不算新奇但要把准确率做到99%以上、还能在不同负载工况下保持泛化中间有一堆细节值得记录。这篇文章我尽量把选型原因、核心原理、数据预处理、模型搭建、踩坑经验一次讲透适合正在做故障诊断、状态监测或者刚接触视觉Transformer想找个实际场景练手的朋友。不管是刚入门的学生还是现场做设备维护的工程师照着这条路线基本都能复现出来。1. 方案选型与整体设计思路1.1 为什么不用传统的特征工程以前做轴承故障诊断主流做法是从振动信号里提取时域指标和频域指标。时域就是均方根值、峰值因子、峭度这些频域就是包络谱、功率谱密度再配合支持向量机或者随机森林做分类。这套方法在单一工况下确实能用换个转速、换个负载特征分布一变模型往往就要重新调。问题出在特征选择太依赖专家经验。一个信号里有转频、谐波、边带、冲击响应到底哪个频带包含故障信息对不同的轴承和不同的故障类型是不一样的。手动挑特征本质上是在用先验知识做压缩丢掉的信息可能恰好是关键信息。所以我后来更倾向于让模型自己从原始信号里学特征尤其是对变工况场景数据驱动的方式鲁棒性明显更好。1.2 一维信号直接分类的局限既然要让模型自己学特征最直接的做法是把一维振动信号直接送入一维卷积网络。确实有不少论文这么做训练也快。但实际做下来发现一个问题轴承故障的冲击成分在原始波形里持续时间很短能量占比小一维卷积虽然能提取局部模式但对某个时刻、某个频段同时出现什么特征这种联合表达建模效率不高。举个例子内圈故障和外圈故障都可能产生冲击但冲击的间隔不同、激励的共振频带不同。这些差异在原始波形里体现得很隐晦而在时频图里一眼就能看出区别——一个呈周期性横条一个在特定频带出现规律亮点。所以先做时频变换相当于把信号信息重新组织成更适合二维视觉模型处理的形式后面再接图像分类模型就会顺利很多。1.3 为什么选SwinTransformer而不是ViT把时频图当作图像来处理最开始大家用的是ResNet这类卷积网络。ResNet的问题在于感受野固定对时频图里那种局部宽频带共振 长周期重复模式的多尺度特征表达能力有限。后来换成Vision TransformerViT全局自注意力理论上能看到整张图但计算复杂度是图像尺寸的平方而且在小样本数据集上很容易过拟合。SwinTransformer的设计恰好卡在这个中间地带它把注意力限制在局部窗口内计算复杂度变成线性增长又通过移位窗口机制实现跨窗口信息交互。再加上它的层级结构类似CNN的Feature Pyramid能同时建模高频冲击细节和低频调制周期非常适合时频图这种同时包含宽带瞬态和窄带稳态成分的数据。所以我最后的方案定为小波时频图 SwinTransformer不是因为它最新而是因为它的归纳偏置和这种数据最匹配。2. 核心原理拆解小波时频图与Swin的关键设计2.1 连续小波变换的数学直觉与参数选择连续小波变换的公式长这样WT(a, b) 1/sqrt(a) * integral(x(t) * psi*((t - b) / a) dt)其中a是尺度参数控制小波函数的伸缩b是平移参数控制小波函数在时间轴上的位置。把小波函数沿着信号滑动并计算内积就得到了二维的系数矩阵横轴是时间纵轴是尺度。这里有个关键直觉尺度a越大小波函数被拉伸得越宽对应的时域窗口越长频率分辨率越高、时间分辨率越低尺度a越小情况反过来。这正好解决了短时傅里叶变换里窗长固定的矛盾。轴承故障产生的冲击是宽频瞬态信号用大尺度看趋势、用小尺度看冲击细节一次变换全都能覆盖。尺度到频率的换算关系是近似公式f Fc * fs / aFc是小波函数的中心频率fs是采样率。实际做的时候我一般用pywt库import pywt import numpy as np scales np.arange(1, 256) coefficients, frequencies pywt.cwt(signal, scales, waveletmorl, sampling_period1/fs)我通常选Morlet小波因为它的时频窗乘积最小频率选择性好而且和轴承冲击响应的形态比较接近。尺度数我常用256或者512太少会损失频率分辨率太多计算变慢但精度提升不明显。做完CWT之后取系数幅度谱再做一次对数压缩把动态范围压下来最后归一化到0到255存成灰度图。2.2 SwinTransformer的关键设计SwinTransformer有几个设计点需要真正理解不然改参数的时候会很盲目。第一是Patch Partition。输入图片先切成4×4的小块每个小块展平映射成向量这一步就把224×224的图变成了56×56个patch每个patch的维度是96。这样等于把图像下采样了4倍后续SwinTransformer就在patch序列上做计算。第二是窗口自注意力。Swin在每一层把patch分成多个7×7的窗口只在窗口内部计算自注意力。相比ViT的全局自注意力计算量大大降低。比如56×56的patch图分窗口后每个窗口49个token整层计算范围仍限制在各窗口内复杂度是窗口数量的线性倍而全局注意力是全部token两两计算。第三是移位窗口。只做局部窗口的话窗口之间没有信息交流感受野就受限了。Swin的处理很巧妙在下一层把窗口整体平移几个patch比如偏移3个patch这样原本窗口边缘的patch在下一层就被分到了不同的窗口等于完成了跨窗口交互。W-MSA和SW-MSA交替使用既保证效率又保证特征能跨区域融合。第四是Patch Merging。每经过几个BlockSwin就把相邻2×2的patch合并成一个特征图尺寸减半、通道数加倍。这个设计让网络具备层级多尺度表征能力浅层看细节、深层看全局处理时频图这种多尺度信号正好合适。3. 数据准备与预处理决定成败的80%3.1 数据集划分与防泄漏很多人在这个环节翻车。一开始我用的公开轴承数据包含正常、内圈故障、外圈故障、滚动体故障等几个大类每类下面又有不同损伤尺寸和负载工况。原始数据是一段很长的连续振动记录需要滑窗切分成多个样本。这里最容易出问题的是数据泄漏。如果先把整段信号按顺序切成900个样本再随机分成训练集和测试集那么同一个连续记录段里的样本可能同时出现在训练集和测试集里。相邻样本高度相关测试集准确率会虚高模型真实泛化能力被严重高估。正确做法是按照记录段来划分比如某段连续信号只进训练集或只进测试集或者按照文件维度划分后再切样本保证测试集中的样本和训练集中的样本来自不同的采集时段。我当时实际是这样处理的以某个转速、某个负载下的原始记录文件为单位先按文件名做分层划分比如80%文件进训练、20%进测试再对每个文件做滑窗采样。这样同一文件下的样本不会横跨两个集合泄漏问题基本杜绝。如果要做不同负载的泛化测试就按负载级别划分训练和测试集比如只用负载A训练用负载B测试更接近实际工业场景。3.2 时频图生成的具体参数振动信号的采样率通常是12kHz或者48kHz。考虑到轴承故障特征频率和共振频带我一般取4096个点作为单个样本的时长这样在12kHz采样率下大概是0.34秒至少覆盖转频的多倍周期。生成时频图的完整流程如下def signal_to_tf_image(signal, fs12000, waveletmorl, scales_num256): scales np.arange(1, scales_num 1) coefficients, _ pywt.cwt(signal, scales, wavelet) amplitude np.abs(coefficients) # 对数压缩拉伸低幅值细节 log_amp np.log1p(amplitude) # 归一化到0-255 normalized (log_amp - log_amp.min()) / (log_amp.max() - log_amp.min() 1e-8) image (normalized * 255).astype(np.uint8) return image这里我要提一个细节时频图到底存灰度图还是彩色图。很多人喜欢用matplotlib的imshow加cmapjet存成彩色图图是好看了但存下来是三通道RGB三个通道内容几乎一样等于白白增加了计算量。我实测下来灰度图单通道输入效果完全够用模型用同样的参数量去学单通道特征收敛更快。当然如果你后续要做迁移学习用ImageNet预训练权重就必须用三通道输入但从零训练的话单通道更划算。3.3 数据增强的边界时频图有一个和普通自然图像完全不同的地方它的纵轴是频率横轴是时间物理含义明确。所以数据增强不能随便做。水平翻转可以相当于把信号在时间轴上反转物理上仍是一个合理的振动信号随机裁剪也问题不大加噪声适度也可以。但垂直翻转绝对不行那等于把高频和低频互换频率轴语义完全混乱模型会学到错误映射。尺度变换也要小心。对时频图做缩放相当于改变信号的频率范围轴承故障特征频率是和转速强相关的随意缩放可能会破坏故障特征和转频的比例关系。我一般只用轻微缩放0.9到1.1倍不做大幅几何变换。每种故障类型我最终生成了800张时频图训练集640张、验证集80张、测试集80张。这个数据量对Swin从小样本角度来说不算宽裕后面会提到怎么靠训练策略补。4. PyTorch模型搭建与训练实战4.1 完整的模型代码结构SwinTransformer的实现可以自己写也可以用timm库里现成的封装。我自己写过一个复现版本也用过现成库整体上推荐后者因为官方实现里有大量经过验证的细节比如相对位置编码、LayerNorm位置、drop path设置自己抄一遍容易出错。下面是用库搭建分类模型的最小代码import torch import torch.nn as nn import timm class SwinClassifier(nn.Module): def __init__(self, num_classes5, pretrainedFalse): super().__init__() self.backbone timm.create_model( swin_tiny_patch4_window7_224, pretrainedpretrained, num_classes0, # 只要特征向量 ) self.classifier nn.Sequential( nn.Linear(self.backbone.num_features, 512), nn.ReLU(inplaceTrue), nn.Dropout(0.3), nn.Linear(512, num_classes), ) def forward(self, x): return self.classifier(self.backbone(x))这里分类头我多加了一层全连接加Dropout。别小看这个Dropout工业小数据集上Swin的特征维度是768维直接映射到5类很容易过拟合中间加一层带Dropout的全连接能明显提升验证集表现。4.2 优化器、学习率与正则策略训练参数我直接给出实测可用的组合优化器: AdamW 初始学习率: 1e-4 权重衰减: 0.05 Batch size: 32 Epochs: 100 学习率调度: CosineAnnealingLR首10个epoch做线性warmup 损失函数: CrossEntropyLoss with label smoothing0.1一个容易被忽略的点是warmup。Swin的Transformer结构对学习率很敏感一上来就全量学习率很容易把训练搞崩。我习惯前10个epoch从1e-6线性升到1e-4后面再用余弦退火慢慢降。实测下来不加warmup时损失会在前几个epoch出现不规律的震荡加了之后训练曲线明显平滑。label smoothing也很关键。轴承分类任务不是完全互斥的比如轻微内圈故障和轻微滚动体故障的特征有重叠硬标签0或1会让模型过度自信。smoothing设为0.1之后测试集的宏平均F1通常能提升1到2个百分点。关于是否使用ImageNet预训练权重我的实验结论是从零训练效果更好。原因很简单时频图和自然图像的分布差异太大ImageNet预训练学到的边缘纹理特征在时频图上意义不大反而需要微调去覆盖大量原有知识。除非你的训练样本极少比如每类只有几十张否则不建议用预训练权重。4.3 训练流程与评估指标训练主循环我写得比较常规但有一个关键点验证集指标和训练集指标一起记录并且只在验证集最优时保存模型权重。不能用测试集做早停否则测试集就变成了模型选择的一部分最终报告出来的指标会偏乐观。我的训练伪代码如下best_f1 0 for epoch in range(total_epochs): train_loss, train_acc train_one_epoch(model, train_loader, optimizer, criterion) val_loss, val_metrics evaluate(model, val_loader) if val_metrics[macro_f1] best_f1: best_f1 val_metrics[macro_f1] torch.save(model.state_dict(), best_swin.pth) scheduler.step()评估指标不能只看准确率。轴承数据通常各故障类别样本数不平衡正常样本总是最多准确率会被多数类拉高。我重点看三个指标宏平均F1、混淆矩阵、单类准确率。宏平均F1对少数类更敏感混淆矩阵能直观看到哪两类经常被搞混——我这边的实验里最常混淆的是内圈故障和滚动体故障因为两者的冲击特征频率在某些转速下比较接近。训练完成后我把测试集时频图输入模型提取倒数第二层特征做t-SNE降维可视化结果里各类别能明显聚成五个团。这一步不能作为定量指标但能帮助确认模型学到的是真正的故障差异而不是数据集划分带来的偶然模式。5. 常见问题与排查技巧实录5.1 时频图明明没问题模型却不收敛如果确认数据划分没问题、标签也没错但loss就是不降或者降得很慢先检查输入归一化。时频图如果每张图单独做min-max归一化不同样本的幅值尺度就被抹平了模型学不到幅值大小这个重要线索。比如正常轴承振动幅值小、故障轴承幅值大这个信息在单独归一化之后会丢失。正确做法是在整个训练集上统计幅值范围统一用全局的min和max做归一化。第二个检查点是学习率。Transformer类模型对学习率的上限很敏感1e-3往往就崩降到1e-4通常能稳住。如果还不收敛把warmup拉长到20个epoch再不行就降低batch size并同步调低学习率。5.2 过拟合严重怎么办Swin在小样本上过拟合是常态。训练集准确率轻松到99%以上验证集却卡在85%这种情况我遇到过很多次。处理顺序我建议这样先加数据增强括随机裁剪、水平翻转、小幅缩放再加大Dropout分类头加到0.5、开启backbone的drop path率0.2然后把weight decay从0.05提到0.1最后考虑换模型规模从base降到tiny模型参数量直接减少好几倍。还有一个最原始但有效的办法多生成样本。用滑窗采样的重叠率从50%提高到75%样本量几乎翻倍。只要确保划分时同一记录段不跨越训练测试集这个办法是安全且免费的。5.3 训练显存和速度怎么优化Swin不是轻量模型224×224输入、batch size 32一张12G显存的卡勉强够用。如果显存溢出优先调小窗口尺寸或者patch size而不是调小batch。另外PyTorch自带的自动混合精度AMP能省下不少显存同时训练速度提升明显但精度几乎没有损失。再讲一个我之前踩过的坑在CPU上生成时频图再存到硬盘生成2500张图花了将近四十分钟每次调一次参数就要重新生成一遍非常痛苦。后来我把生成逻辑改成在线生成数据加载时实时做CWT虽然训练时每次都要算一遍小波变换但因为可以用多进程并行总体性价比反而更高。不过在线生成的缺点是会增加每个epoch的训练时间所以更推荐的做法是先离线生成好全部时频图但如果要频繁改CWT参数就改成在线生成。5.4 分类效果不错但泛化测试崩了还有一个很隐蔽的问题模型在随机划分的测试集上表现不错但换成不同负载下的数据后准确率暴跌。这说明模型学到了和负载相关的分布特征而非纯粹的故障特征。这种场景下可以考虑两种对策。第一种是数据层面把多个负载的数据混合进训练集并在划分时明确区分负载第二种是模型层面用领域对抗训练或者把频谱做归一化比如去除转频分量让模型不过度依赖绝对幅值。后者实现复杂度高我通常先试第一种成本最低且效果立竿见影。我个人的体会是这套小波时频图 SwinTransformer的组合最适合的是对设备整体运行状态做自动分类和初期报警而不是替代现场专家做精细化故障程度判断。真正放到产线上环境噪声、转速波动、信号衰减都会带来新的挑战但时频图加局部注意力这个思路目前看是最能兼顾精度和工程可行性的方向之一。做工业级项目时建议再配合信号处理领域的一些经典方法做交叉验证双保险总是更稳妥。