在旋转机械故障诊断这个方向上，轴承的振动信号分析永远是绕不开的话题。我之前做过几轮基于传统特征提取的项目，也试过直接喂一维信号给卷积分网络，最终在实际项目里稳定复现、效果也最让我放心的方案，是小波时频图加SwinTransformer这套组合。简单说，就是先把一维振动信号做连续小波变换，变成二维时频图，再交给SwinTransformer做分类。整个流程基于Python和PyTorch实现，代码结构清晰，训练效率和识别精度都不错。

这个思路不算新奇，但要把准确率做到99%以上、还能在不同负载工况下保持泛化，中间有一堆细节值得记录。这篇文章我尽量把选型原因、核心原理、数据预处理、模型搭建、踩坑经验一次讲透，适合正在做故障诊断、状态监测，或者刚接触视觉Transformer想找个实际场景练手的朋友。不管是刚入门的学生，还是现场做设备维护的工程师，照着这条路线基本都能复现出来。

1. 方案选型与整体设计思路

1.1 为什么不用传统的特征工程

以前做轴承故障诊断，主流做法是从振动信号里提取时域指标和频域指标。时域就是均方根值、峰值因子、峭度这些，频域就是包络谱、功率谱密度，再配合支持向量机或者随机森林做分类。这套方法在单一工况下确实能用，换个转速、换个负载，特征分布一变，模型往往就要重新调。

问题出在特征选择太依赖专家经验。一个信号里有转频、谐波、边带、冲击响应，到底哪个频带包含故障信息，对不同的轴承和不同的故障类型是不一样的。手动挑特征，本质上是在用先验知识做压缩，丢掉的信息可能恰好是关键信息。所以我后来更倾向于让模型自己从原始信号里学特征，尤其是对变工况场景，数据驱动的方式鲁棒性明显更好。

1.2 一维信号直接分类的局限

既然要让模型自己学特征，最直接的做法是把一维振动信号直接送入一维卷积网络。确实有不少论文这么做，训练也快。但实际做下来发现一个问题：轴承故障的冲击成分在原始波形里持续时间很短，能量占比小，一维卷积虽然能提取局部模式，但对"某个时刻、某个频段同时出现什么特征"这种联合表达，建模效率不高。

举个例子，内圈故障和外圈故障都可能产生冲击，但冲击的间隔不同、激励的共振频带不同。这些差异在原始波形里体现得很隐晦，而在时频图里一眼就能看出区别——一个呈周期性横条，一个在特定频带出现规律亮点。所以先做时频变换，相当于把信号信息重新组织成更适合二维视觉模型处理的形式，后面再接图像分类模型就会顺利很多。

1.3 为什么选SwinTransformer而不是ViT

把时频图当作图像来处理，最开始大家用的是ResNet这类卷积网络。ResNet的问题在于感受野固定，对时频图里那种"局部宽频带共振 + 长周期重复模式"的多尺度特征，表达能力有限。

后来换成Vision Transformer（ViT），全局自注意力理论上能看到整张图，但计算复杂度是图像尺寸的平方，而且在小样本数据集上很容易过拟合。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 / a

Fc是小波函数的中心频率，fs是采样率。实际做的时候我一般用pywt库：

import pywt import numpy as np scales = np.arange(1, 256) coefficients, frequencies = pywt.cwt(signal, scales, wavelet='morl', sampling_period=1/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。每经过几个Block，Swin就把相邻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, fs=12000, wavelet='morl', scales_num=256): 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加cmap='jet'存成彩色图，图是好看了，但存下来是三通道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_classes=5, pretrained=False): super().__init__() self.backbone = timm.create_model( 'swin_tiny_patch4_window7_224', pretrained=pretrained, num_classes=0, # 只要特征向量 ) self.classifier = nn.Sequential( nn.Linear(self.backbone.num_features, 512), nn.ReLU(inplace=True), 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 smoothing=0.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"的组合，最适合的是对设备整体运行状态做自动分类和初期报警，而不是替代现场专家做精细化故障程度判断。真正放到产线上，环境噪声、转速波动、信号衰减都会带来新的挑战，但时频图加局部注意力这个思路，目前看是最能兼顾精度和工程可行性的方向之一。做工业级项目时，建议再配合信号处理领域的一些经典方法做交叉验证，双保险总是更稳妥。