基于连续小波变换与SwinTransformer的轴承故障诊断实践

国产轴承故障诊断这块，这些年发文章、搞落地的思路已经比较统一了：信号处理负责把原始振动数据“翻译”成模型能看懂的特征，深度学习模型负责在这些特征里自动挖规律。这个项目走的正是这一套组合拳里的高配路线——用连续小波变换把一维振动信号变成二维时频图，再喂给SwinTransformer做故障分类。我实际跑下来，这套方案的识别精度和泛化能力确实比传统CNN路子要稳一截，但里面的坑也不少。这篇就把我的完整实操过程、选型逻辑和踩过的坑都梳理一遍，给打算做同类工作的朋友提供个参考。

1. 项目整体设计与思路拆解

1.1 为什么落在“轴承故障诊断”这个点上

旋转机械里，轴承是最容易出问题的部件之一，大量设备故障的根源都能追溯到轴承磨损或局部损伤。故障模式通常分四类：内圈故障、外圈故障、滚动体故障、保持架故障，加上正常状态一共五类。这些故障在振动信号上有比较明显的冲击特征，但现场工况复杂，负载、转速一变化，特征形态就会漂移，单纯靠人工经验看波形定故障，效率低且容易漏判。所以把这个场景作为深度学习诊断的切入点是合理的：数据获取相对容易，故障类别明确，而且诊断结果的准确性可以直接影响设备维护决策，工程价值很高。

我用的数据是某高校公开的滚动轴承振动数据集，采样频率12kHz，包含正常、内圈故障、外圈故障、滚动体故障等多种状态，每种状态又有不同故障直径。这类公开数据最大的好处是方便横向对比论文结果，但也要注意它数据分布相对干净，和现场实测数据还是有差距，后期落地时要专门处理噪声和工况迁移问题。

1.2 特征层面：为什么看时频图而不是原始波形或频谱

这是整个项目最关键的一个选型。直接拿原始振动波形喂CNN，效果往往不太理想，因为一维信号里故障冲击的相位和幅值变化对平移、噪声太敏感，模型要学的模式非常不稳定。做FFT频谱虽然能看到故障特征频率，但谱是全局平均的结果，把“什么时候出现冲击”这个时间信息完全抹掉了，而很多故障的差异恰恰体现在瞬态冲击的时变规律上。

小波时频图补偿了这两边短板。它把一维时间序列映射成“时间-频率-幅值”的三维表达，本质是给信号做一系列带通滤波，每个尺度对应一个频带，输出是信号沿着时间轴在各频带的能量分布。故障冲击在时频图上会表现为特定频率附近的亮色条带，不同故障位置、不同损伤程度的纹理特征有明显差异，非常适合图像模型学习。

对比一下三种常见做法：

短时傅里叶变换（STFT）窗口一固定，高频段时间分辨率不足、低频段频率分辨率不足的问题就来了。而小波变换的窗口是自适应的：尺度小（对应高频）时窗口窄、时间定位准；尺度大（对应低频）时窗口宽、频率定位准。这种性质对轴承冲击信号这种宽频瞬态成分特别友好。

1.3 模型层面：为什么用SwinTransformer而不是CNN或普通ViT

拿到时频图后，接下来的任务就是分类。传统做法是搭个ResNet之类的CNN，卷积核堆叠扩大感受野，能学出不错的特征。但CNN有个结构性限制：全局上下文信息需要很多层才能逐步传递到高层特征里，对时频图中故障特征分散在多个频带、多个时间段的情况，建模效率一般。

ViT把图像切成patch后用全局自注意力建模，理论上能直接捕捉任意两个位置的依赖关系。但普通ViT有两个问题：计算复杂度随图像尺寸平方增长，训练时需要海量数据，用在几百上千张样本的故障数据集上很容易过拟合。

SwinTransformer的折中方案是这套架构能落地的核心：自注意力被限制在固定大小的窗口内（比如7×7 patch范围），窗口内部建模局部关系，窗口之间通过“移位窗口”机制在相邻层交换信息，从而在保持近似全局建模能力的同时把计算量降到了线性级别。另外窗口化的结构天然保留了卷积的局部先验，收敛性和泛化性都比全局ViT好。实际训练时，Swin在小规模工业数据集上也能快速收敛，这一点很对故障诊断这个场景的胃口。

2. 小波时频图生成实战

2.1 小波变换核心原理速通

我不打算堆公式劝退，但几个关键概念还是得说清楚。连续小波变换（CWT）做的事情，是把一个母小波函数在不同尺度a下做伸缩、在时间轴上平移，然后和信号做内积：

尺度a决定母小波的伸缩程度：a小、波形被压缩，对应高频分析；a大、波形被拉伸，对应低频分析。对轴承振动信号，我们需要覆盖几百赫兹到几千赫兹的频带范围，所以尺度的起点和终点要根据采样频率和感兴趣频段来定。

尺度与频率的对应关系近似为：f = fs × fc / a，其中fs是采样频率，fc是母小波的中心频率。也就是说，选定尺度序列，就等价于选定了一组从高频到低频的分析频带。

2.2 参数怎么选：小波基、尺度序列与颜色映射

小波基的选择直接决定时频图质量。实测下来，复Morlet小波在轴承振动信号上的效果最稳。原因是Morlet小波本质上是一个被高斯窗包络的正弦波，和轴承故障冲击波形（衰减的正弦振荡）形态接近，匹配度高，时频能量更集中。而像哈尔小波这类不连续的基，得到的时频分布比较碎，分类效果明显差一些。

尺度序列范围取决于你关心的频带。采样频率12kHz，奈奎斯特频率是6kHz，轴承故障的固有振动频带通常在几百赫兹到4000赫兹之间。我用的尺度范围是32到512，对数均匀分布生成128个尺度。这个设置基本覆盖了主要故障频带，又不至于让低频频段占过多冗余。亲测在128个尺度下生成的时频图，224×224输入下纹理够清楚，再加密尺度对精度提升有限，但计算时间会显著增加。

代码实现用pywt就很方便：

python复制import pywt
import numpy as np
from scipy.signal import resample

def generate_cwt_spectrogram(signal, fs=12000, scales_num=128):
    # 生成尺度序列：对数均匀分布
    scales = np.logspace(np.log2(32), np.log2(512), scales_num, base=2)
    # CWT计算
    coefficients, freqs = pywt.cwt(signal, scales, wavelet='morl', sampling_period=1/fs)
    # 取模得到幅值时频图
    spectrogram = np.abs(coefficients)
    # 对数变换压缩动态范围，突出弱冲击成分
    spectrogram = np.log1p(spectrogram)
    # 归一化到[0,1]区间
    spectrogram = (spectrogram - spectrogram.min()) / (spectrogram.max() - spectrogram.min() + 1e-8)
    return spectrogram, freqs

注意：时频图矩阵的方向是（尺度数, 时间点数），和图像行列正好对应。保存或送入模型前，这里的转置和对齐要注意，不然图像内容就反了。

2.3 从时频图到训练集的“工业化”细节

生成时频图只是第一步，要把几百段原始信号变成几千张训练图片，中间有几个细节直接决定模型效果：

第一是样本切片方式。原始振动信号通常很长（几十万点），直接整段生成时频图再切图，信息冗余且样本量不够。标准做法是把信号切成固定长度的片段（我用的每段2048个点，约0.17秒），每个片段生成一张时频图。注意相邻片段之间可以设置重叠，我用50%重叠率，把小样本数据集扩充了近一倍。但务必先切片段再生成时频图，不要先生成整段时频图再在图像上裁剪——后者会引入时间上的混叠信息，标签对应关系会变得不干净。

第二是尺寸对齐。SwinTransformer需要固定输入尺寸，我统一把时频图resize到224×224。resize用PIL的双三次插值，别用OpenCV默认的双线性，后者在缩小图片时容易把高频细节抹掉。

python复制from PIL import Image

def prepare_image(spectrogram):
    # 转成uint8灰度图
    img = (spectrogram * 255).astype(np.uint8)
    img = Image.fromarray(img, mode='L')
    # 统一尺寸
    img = img.resize((224, 224), Image.BICUBIC)
    arr = np.array(img, dtype=np.float32) / 255.
    # 扩展成3通道，保持和ImageNet预训练输入结构一致
    arr_rgb = np.stack([arr, arr, arr], axis=-1)
    return arr_rgb

有人会问：时频图是单通道灰度图，为什么要复制成3通道？因为后面要用ImageNet上预训练好的SwinTransformer做迁移学习，模型第一层卷积（或者patch embedding层）的设计是按RGB三通道来的，直接跑灰度图会维度不匹配。复制成3通道不增加信息量，但能顺利加载预训练权重，省去从头训一大截时间。

第三是数据增强策略。时频图不是普通自然图像，时间和频率两个方向都有物理意义——时间轴翻转等于把信号倒着放，频率轴本来是高频在上还是低频在上取决于绘图方式，增删这些都会改变物理含义。我实测下来，随机水平翻转（时间方向）和轻微随机裁剪（10%以内）是可以用的，增益比较明显；但垂直翻转（频率方向翻转）会严重破坏频谱的物理意义——低频变高频、高频变低频，模型学到的是错误映射，精度明显下降且难以解释。因此数据增强要克制，尤其是不要用带物理语义的强增强。

3. SwinTransformer模型搭建与关键解析

3.1 Swin的核心组件拆解

SwinTransformer整体可以看成一个“金字塔结构”的Transformer编码器，核心组件有三个：

Patch Embedding（Patch切分与嵌入）是把图像切分成4×4的小块，每个小块展平成向量后通过一个线性层映射到embedding维度。224×224的输入，经过4×4切分后得到56×56个patch，每个patch的向量维度设为96（Swin-Tiny的配置）。这个操作等效于一个kernel_size=4, stride=4的卷积，实际代码里就是这样实现的。

Swin Transformer Block负责特征提取，每个Block内部由两部分组成：第一部分是窗口多头自注意力（W-MSA），只在固定的7×7窗口内做自注意力计算；第二部分是移位窗口多头自注意力（SW-MSA），窗口位置整体平移（偏移3个patch），让信息可以跨窗口流动。W-MSA负责局部建模，SW-MSA负责全局连接，两个模块交替堆叠，就构成了一个兼顾局部和全局的双层结构。

Patch Merging（下采样合并）跟CNN里的池化层作用类似。每经过一个阶段，把2×2相邻的patch合并成1个，通道数翻倍。这个设计让网络能像CNN一样堆出多尺度特征：浅层分辨率高、通道少，捕捉细节纹理；深层分辨率低、通道多，凝聚语义信息。

值得强调的是，Swin的窗口注意力里有一个容易被忽略的细节——相对位置偏置表（relative position bias）。每个查询和键的相对位置映射到一个可学习的偏置值，加到注意力得分上。这个偏置让模型天然具备平移等变性，也正是它让Swin在没有额外数据增强的情况下比ViT更容易收敛。我在训练中尝试过去掉这个偏置，验证集准确率掉了约3个百分点，证明它对小数据集训练是实打实有帮助的。

3.2 Pytorch实现核心代码

社区里有多个SwinTransformer的开源Pytorch实现，我建议直接用官方的timm库版本，而不是自己从零写，因为注意力可视化、预训练权重导入这些问题timm都替你处理好了。关键代码可以浓缩成几十行：

python复制import timm
import torch.nn as nn

class BearingDiagSwin(nn.Module):
    def __init__(self, num_classes=5, pretrained=True):
        super().__init__()
        # 使用timm官方的Swin-Tiny结构
        self.backbone = timm.create_model(
            'swin_tiny_patch4_window7_224',
            pretrained=pretrained,
            num_classes=0,          # 去掉原始分类头，只拿特征
        )
        # 自定义分类头
        self.classifier = nn.Sequential(
            nn.Linear(768, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(256, num_classes),
        )

    def forward(self, x):
        feat = self.backbone(x)      # 输出维度 (batch, 768)
        out = self.classifier(feat)
        return out

这一段代码就够搭出完整的模型结构。这里的768是Swin-Tiny经过四层stage后最终输出的特征维数。如果样本量特别少，也可以把backbone参数冻住，只训练分类头，跑通流程后再解冻微调。

3.3 针对时频图的模型配置调整策略

有人会直接拿Swin原始配置放在时频图数据集上跑，效果反而一般，原因是224×224的输入对时频图来说，patch切分（4×4）后丢失了不少窄带细节。我的经验是两条路：

一条是保持224输入，但把第一阶段的embedding维度从96降到64，减少小数据集下的参数量。Swin-Tiny全参大约是2800万，工业小数据集上容易过拟合，减到64维可以有效缓解。代价是预训练权重没法直接加载，需要从头训练，训练时间相应拉长。

另一条是保留timm预训练权重，但输入分辨率改成192或256。Swin的窗口划分是基于7×7窗口的，192能被7整除，224也能，256不行（会余5），强行用256会触发窗口padding分支，推理速度变慢的同时精度未必提升。实测224是最合适的。这里说句题外话：用Swin如果输入尺寸选不好，代码会自动走F.pad分支，显存占用直接上升，这部分开销完全没必要。

最终我采用的配置是：Swin-Tiny结构、输入224×224×3、patch_size=4、window_size=7、embedding_dim=96、四个stage的block深度分别是2/2/2/2（比原始2/2/6/2浅，减少过拟合）、head数量3/6/12/24。参数量约1500万，比原版少近4成，在小样本轴承数据上收敛更稳。

4. 训练全流程与关键评估实验

4.1 数据集划分：最容易出假高分的地方

这个环节我必须单拎出来说，因为见过太多人在这里“不经意”地刷高了指标。轴承振动信号是时间序列，同一段信号经滑动窗口切出来的相邻样本，内容高度相关。如果划分训练集和测试集时混入了来自同一原始信号段的样本，模型等于提前见过答案，准确率虚高是必然的。

正确做法是先把原始信号分文件按类别划分，确保同一段物理信号的样本全部落在一个集合里，再进行二次划分。我的做法是：按8:1:1比例先划分原始信号文件，再在文件内部做滑动窗口切样本。这样训练集、验证集、测试集的样本之间不存在重叠来源，评估结果才有参考意义。

分类别统计样本量，每个类别大约400到600张时频图不等，存在轻度类别不平衡。训练时用WeightedRandomSampler对样本加权，避免模型偏向样本量多的类别。

4.2 训练超参与优化配置

超参这块，我实验出来的组合如下：

训练代码：

python复制import torch
from torch.utils.data import DataLoader
from timm.scheduler import CosineLRScheduler

model = BearingDiagSwin(num_classes=5, pretrained=True).cuda()
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05)
scheduler = CosineLRScheduler(optimizer, t_initial=80, lr_min=1e-6)

for epoch in range(80):
    model.train()
    total_loss, correct, total = 0, 0, 0
    for x, y in train_loader:
        x, y = x.cuda(), y.cuda()
        optimizer.zero_grad()
        out = model(x)
        loss = criterion(out, y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
        optimizer.step()
        total_loss += loss.item() * x.size(0)
        correct += (out.argmax(1) == y).sum().item()
        total += y.size(0)
    scheduler.step(epoch)
    print(f'epoch {epoch+1}: loss {total_loss/total:.4f}, acc {correct/total:.4f}')

梯度裁剪这个操作容易被忽略，但对Transformer很关键：注意力机制对梯度爆炸敏感，不加clip有时候前几个epoch就loss飞掉。max_norm=5.0是我调出来的稳当值。

4.3 评估指标与结果分析

只看整体准确率太单薄，故障诊断场景我更看重每个类别的精确率和召回率，因为不同故障类型在时频图上的区分难度差异很大。测试集上的最终结果：

整体准确率96.9%。滚动体和保持架故障识别率偏低属于意料之中：这两类故障产生的冲击能量相对内圈和外圈弱不少，而且在时频图上的特征容易被噪声淹没。这一点也给后续优化指明方向——对这两类做专门的数据增强或者设计更细粒度的特征分支。

混淆矩阵里，误判集中发生在滚动体故障被误判为正常状态、保持架故障被误判为内圈故障这两处。前者说明冲击特征不够明显时模型倾向保守；后者说明两者的频带分布存在部分重叠。这个信息只有看混淆矩阵才能拿到，只看总体准确率是看不出来的。

对比实验我做了三组：CNN基线（ResNet18+时频图）、纯ViT（DeiT-Small+时频图）和本方案（Swin-T+时频图）。结果Swin比ResNet高约2.8个百分点，比DeiT高约1.5个百分点。ResNet差在最开始几层卷积的局部感受野对宽频噪声敏感；DeiT则是在小样本下收敛不够充分，出现了过拟合迹象。

5. 常见问题与排查技巧实录

5.1 时频图生成效率低

小波变换本质是内积运算，128个尺度对2048点信号做一遍，单个样本在CPU上大约要0.2秒。几千个样本还好，但如果做实时数据流或样本量到几万张，计算时间就比较可观了。处理方式是预生成缓存：先离线把全量样本的时频图保存为.npy或图片文件，训练时直接读图，不再实时计算，训练速度能提升一个数量级。

5.2 模型Loss不下降或直接NaN

先检查输入数据：时频图有没有做归一化？归一化方式是不是把零值压成了有效值，导致log1p之后出现NaN？然后检查标签：类别编号是0到4还是从1开始？很多分类代码里标签编号不连续，会导致CrossEntropyLoss内部Assert报错或梯度异常。最后检查学习率：用ImageNet预训练权重时初始学习率超过1e-3基本必炸。

5.3 出现极高准确率但肯定是假象

我遇到过一位同行跑出99%+“准召率”，后来查代码发现滑动窗口只做了0.05秒、重叠率95%，导致同源信号被无限复制，训练集和测试集高度重叠。排查手段很简单：训练集和测试集各抽一批样本，然后跨集合计算逐样本相似度。时频图是确定的数学变换，相同输入必然得到相同输出，相似度高就说明数据切分不干净。修正我上面说的“按信号文件划分”方案后，准确率降到了正常范围（96%左右），这才是真实水平。

5.4 显存不足

Swin-Tiny在224×224输入、batch=32的情况下显存占用大约6GB，8G显存的显卡其实能跑。如果仍报OOM，优先把batch减半并用梯度累积弥补：每积累两步再更新一次参数，效果基本等效于原batch大小。另外别开cudnn.benchmark的自动调优，它的显存探测阶段在某些卡上会多占几百MB。

5.5 过拟合的常见信号与应对

训练集准确率很快冲到99%，但验证集在85%附近徘徊，这是最典型的过拟合信号。优先检查数据增强是否够用、LabelSmoothing是否打开、分类头Dropout是否在0.2以上。如果还压不住，再考虑用轻量级正则——在损失函数里加上特征向量Frobenius范数正则项，这个技巧在Swin这类Transformer上比传统权重衰减更有效，相当于约束特征空间的范数边界。还有一种更省事的路线：把预训练权重冻结、只训练分类头，先把验证集准确率拉到90%附近，再解冻骨干网络、用更低学习率微调。

5.6 时频图翻转到底能不能做增强

上文提到水平翻转可行的依据是：轴承旋转的启停方向在特征层面不会改变故障类型，时频谱的时间轴翻转只是改变了冲击出现的先后顺序，对频带分布影响很小，模型学到的仍是相对稳定的频带纹理。但垂直翻转不可行，因为频率轴的物理意义在这个任务里直接绑定故障类别。总结一句话：做带物理语义的图像的增强策略前，先想象一下增强后的图还能不能代表原故障的物理本质，不能就不加。

实操心得与后续扩展建议

这个项目做下来，我最深刻的体会是：模型选型不是决定诊断精度的唯一因素，甚至是相对次要的因素。更关键的是把“信号特征表达”这件事吃透——小波基、尺度范围、切片方式，每一个参数的微小变化，对最终准确率的影响都可能超过换一个更强的模型结构。Swin在这里只是恰好匹配了时频图的特征形态，真正拉开差距的是前面的信号处理环节做得够不够细。

对想复现或者扩展这套方案的朋友，我的建议是先从公开数据集跑通全流程，再逐步添加噪声、变工况迁移等更接近实际现场的条件。变工况下模型性能会明显下降，此时可以考虑加入域适配模块或采用对比学习做信号特征对齐。实际设备上做在线诊断时，建议先运行异常检测，再对异常段单独生成时频图做细粒度分类，算力开销比全时段跑Swin低一个数量级，工程上更容易接受。

最后分享一个我在部署阶段学到的血泪经验：现场采集的数据采样率、传感器灵敏度很可能和训练数据集不同，直接套模型往往失灵。我的做法是给模型加一个自动增益归一化层，在输入时频图前先对每个样本做以峰值为基准的最大值归一化，再叠加一个可学习的缩放参数，让网络自己适应不同量纲的输入。这层改动对跨设备部署的帮助非常直接——多台设备共享同一个模型、无需各自微调，节省的可不只是重新训练的时间。