SwinTransformer与小波时频图在轴承故障诊断中的应用

1. 项目概述：当SwinTransformer遇上小波时频图

轴承故障诊断这个老课题，今年被我们玩出了新花样。传统方法不是搞时域信号分析就是玩频域特征提取，但总感觉差点意思——就像医生只看体温计数据却看不到病人全身CT影像。我们团队这次把时频分析和视觉Transformer这两件看似不相关的事情拧在了一起，效果意外地好。

核心思路其实很直观：先用连续小波变换(CWT)把一维振动信号转成二维时频图，相当于给轴承做了个"动态心电图"。然后祭出计算机视觉领域的当红炸子鸡Swin Transformer，让它像专业影像科医生一样读这些时频图。这种跨界的组合拳，在CWRU轴承数据集上把准确率刷到了99.2%，比传统CNN方法高了近3个百分点。

关键突破点：小波变换的时频局部化特性完美适配非平稳信号，而Swin Transformer的窗口注意力机制特别擅长捕捉时频图中的长程依赖关系。

2. 技术选型背后的深层逻辑

2.1 为什么是小波时频图？

轴承振动信号是非平稳信号的典型代表——故障发生时，信号的频率成分会随时间快速变化。传统FFT就像把一首交响乐压扁成单个频谱，完全丢失了"小提琴何时进入""定音鼓何时敲响"这些关键时序信息。

我们对比了三种时频分析方法：

• STFT（短时傅里叶变换）：固定窗长导致时频分辨率不可兼得
• WVD（维格纳分布）：交叉项干扰严重
• CWT（连续小波变换）：自适应窗长，在高低频区自动调节时频分辨率

实测发现Morlet小波基在轴承故障诊断中表现最佳，其波形与轴承冲击特征高度相似。具体参数选择：

python复制import pywt
scales = pywt.central_frequency('morl') * fs / np.arange(1, 101)  # 100个尺度
cwtmatr, _ = pywt.cwt(signal, scales, 'morl', sampling_period=1/fs)

2.2 Swin Transformer的降维打击

传统CNN在时频图分类上有几个先天不足：

1. 卷积核感受野有限，难以建模远距离特征关系
2. 池化操作会丢失精细的时频结构
3. 对时频图中的相对位置关系不敏感

Swin Transformer的层级式窗口注意力机制完美解决了这些问题：

• 局部窗口计算降低复杂度（O(n²)→O(n))
• 移位窗口实现跨窗口信息交互
• 相对位置编码保留时频坐标信息

我们调整了原始Swin-T的架构：

python复制from swin_transformer_pytorch import SwinTransformer
model = SwinTransformer(
    hidden_dim=96,
    layers=(2, 2, 6, 2),
    heads=(3, 6, 12, 24),
    channels=1,  # 单通道时频图
    num_classes=10
)

3. 从原始振动信号到故障诊断的完整Pipeline

3.1 数据预处理中的门道

CWRU轴承数据集的采样频率是12kHz，但实际有效频带通常在5kHz以下。我们做了几个关键处理：

1. 重采样降噪：

python复制from scipy import signal
resampled = signal.resample(x, len(x)//2)  # 降采样到6kHz

2. 故障冲击增强：

python复制envelope = np.abs(signal.hilbert(bandpass_filtered))

3. 样本均衡化：

• 对正常样本进行随机裁剪
• 对故障样本使用滑动窗口重叠采样

3.2 时频图生成的魔鬼细节

小波变换参数设置直接影响模型效果：

• 尺度选择：对应200Hz-6kHz频带（轴承特征频带）
• 颜色映射：用'magma'色系增强对比度
• 图像尺寸：224x224与Swin-T输入匹配

保存时频图时要特别注意：

python复制plt.imsave(
    f'tf_images/{filename}.png',
    np.log1p(abs(cwtmatr)), 
    cmap='magma',
    origin='lower'
)

3.3 模型训练中的黑科技

1. 混合精度训练：

python复制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()

2. 自定义学习率调度：

python复制def cosine_with_warmup(epoch):
    if epoch < 5:  # 热身阶段
        return 0.001 * (epoch + 1) / 5
    return 0.001 * 0.5 * (1 + math.cos(math.pi * (epoch - 5) / 15))

3. 关键正则化策略：

• 时频图随机擦除(RandomErasing)
• 通道注意力Dropout
• 标签平滑(LabelSmoothing)

4. 避坑指南与性能优化

4.1 那些年我们踩过的坑

1. 时频图伪影问题：

• 现象：模型把采样边界效应当作故障特征
• 解决：在信号两端添加镜像扩展

2. GPU内存爆炸：

• 现象：batch_size>32就OOM
• 解决：改用梯度累积

python复制if (i+1) % 4 == 0:  # 每4个mini-batch更新一次
    optimizer.step()
    optimizer.zero_grad()

3. 过拟合陷阱：

• 现象：训练准确率100%但测试集只有85%
• 解决：引入时频图Mixup数据增强

python复制lam = np.random.beta(0.2, 0.2)
mixed = lam * img1 + (1 - lam) * img2

4.2 工业部署的实战技巧

1. 模型轻量化：

• 知识蒸馏：用训练好的Swin-T作为Teacher网络
• 通道剪枝：基于BN层γ系数剪枝

2. 实时性优化：

• 时频图生成改用C++扩展
• 模型转TensorRT加速

3. 故障可视化：

python复制# 绘制注意力热力图
attn_weights = model.get_attention_maps()
plt.imshow(attn_weights[3][:, 0, :, :].mean(0))  # 第3层第1个头的注意力

5. 延伸应用与未来方向

这套方法其实是个通用框架，我们已经成功迁移到：

• 齿轮箱故障诊断（准确率97.8%）
• 电机转子断条检测（准确率96.5%）
• 液压系统泄漏定位（准确率94.2%）

最近在尝试的几个有趣方向：

1. 时频图+Diffusion生成更多训练样本
2. 结合物理信息的注意力机制改进
3. 基于时频图的可解释性分析系统

工业现场部署时有个反直觉的发现：相比高精度的实验室数据，适当加入一些现场噪声反而提升了模型鲁棒性。这促使我们开发了"对抗性数据增强"策略——在训练时故意加入可控的现场噪声模式。