做轴承故障诊断这个方向有一段时间了，踩了不少坑，也积累了一些经验。最近把整套流程重新梳理了一遍，从原始振动信号到最后模型输出，决定用一篇文章把"小波时频图 + SwinTransformer"这套组合方案的完整实现细节记录下来。这篇文章不是论文复现，也不是官方文档翻译，而是基于我在实际项目中反复调参、调整方案后沉淀下来的实操总结。适合已经接触过Python和Pytorch基础、想做故障诊断但不知道从哪下手的同学，也适合已经在用CNN做诊断、想试试Transformer这条路线的工程师。

轴承故障诊断最核心的问题不是"用什么网络"，而是"怎么从振动信号里挖出有效的故障特征"。很多做深度学习的同行一上来就想着堆模型，结果发现在现场数据上精度上不去，问题恰恰出在数据预处理和特征表达上。这篇文章会详细说明为什么选用小波时频图作为模型输入、SwinTransformer相比传统CNN的优势在哪里、整个Pytorch实现过程中需要注意的细节，以及我在实际训练中遇到的几个容易让人崩溃的坑。

1. 为什么是"小波时频图 + SwinTransformer"这对组合

1.1 传统方法的瓶颈在哪里

先说说轴承故障诊断这个任务本身。滚动轴承在运行过程中，振动信号会携带大量状态信息。当轴承出现内圈、外圈、滚动体或保持架故障时，振动信号会产生周期性的冲击成分，这些冲击在时域波形上表现为一个个尖峰，在频域上表现为特定频率处出现边带和峰值。但问题在于，工业现场的振动信号远没有教科书那么干净。

首先是信号的非平稳性。轴承转速波动、负载变化、噪声干扰，都会让振动信号的统计特性随时间变化。频域分析（比如FFT）只能告诉你整个时间窗口内有哪些频率成分，但无法告诉你这些频率成分是什么时候出现的。时域分析虽然保留了时间信息，但很难直观区分不同故障类型。这就是为什么单靠时域特征或频域特征做诊断，泛化能力一直上不去的根本原因。

其次是故障特征的微弱性。早期故障的冲击能量非常小，经常淹没在背景噪声和齿轮啮合振动中。实测试过用原始时域波形直接训练，模型准确率能到90%以上，但现场一换工况就立刻掉到70%以下，这就是典型的过拟合了数据集本身，而不是学到了真正的故障模式。

1.2 时频分析方法怎么选

时频分析是解决非平稳信号问题的经典思路，它把一维的振动信号映射成二维的时频图像，同时保留时间和频率信息。常用的方法包括短时傅里叶变换（STFT）、Wigner-Ville分布和小波变换。我在项目中把STFT和小波变换都试了一遍，简单对比一下。

STFT的思路是给信号加一个固定长度的窗，假设窗内信号是平稳的，然后做FFT，窗沿着时间轴滑动就得到一张时频图。它的硬伤在于窗长一旦确定，时间分辨率和频率分辨率就锁死了，二者互斥。想看清楚低频段，需要长窗，但长窗在时间上的定位能力就变差了。轴承故障信号的特点是瞬态冲击，频率成分在短时间内容易突变，STFT在这类信号上的表达能力其实一般。

小波变换则完全不同。它用一组可伸缩、可平移的小波基函数去匹配信号的局部特征。高频段用窄窗、低频段用宽窗，自动适应信号在不同频段的特性。对于轴承故障这种"高频冲击叠加低频调制"的信号，小波变换的理论适配性明显更好。我当时直接用连续小波变换（CWT）生成的时频图作为输入，可视化效果比STFT清晰得多，故障冲击的频带特征一眼就能分辨。

这里给个表格直观对比：

1.3 SwinTransformer 比 CNN 强在哪

输入图像确定之后，接下来就是选模型。近两年做故障诊断的主流方案依然以ResNet、VGG这类CNN为主，但我在实践中的感受是，CNN在时频图分类上存在两个天然短板。

一个是感受野受限。CNN靠卷积核堆叠来扩大感受野，但每层的感受野始终是局部的，对图像中"远距离依赖关系"的建模能力不足。时频图上的故障特征往往不是孤立的一个点，而是横跨一定时间宽度和频率宽度的条带状区域，并且不同频带之间的关联可能相隔较远，CNN需要非常深的网络才能捕捉到这种长程关系。

另一个是层次化特征的自适应能力。CNN的池化操作是固定的下采样方式，不论图像内容如何，都机械地降低分辨率。SwinTransformer 引入了移动窗口注意力机制，在窗口内部计算自注意力，不同层之间通过窗口移位实现跨窗口信息交互。这样既保留了Transformer全局建模的能力，又通过窗口化大幅降低了计算复杂度，让模型的注意力能够自适应地聚焦在关键故障区域。

我实测下来，同一个数据集上，SwinTransformer（Swin-Tiny级别）的准确率比ResNet50高出2到3个百分点，而且在小样本场景下优势更明显。它也不是完全没有代价——训练时间确实比CNN长，显存占用也更高，但为了精度，这些成本值得。

2. 数据流水线：从原始振动波形到可学习的时频图

2.1 数据来源与信号预处理

我们这次使用的是公开的滚动轴承故障数据集，包含正常状态、内圈故障、外圈故障、滚动体故障四类样本，每个类别又有不同损伤程度。数据采集时的采样率是12kHz，转频范围在几十到几百赫兹，故障特征频率集中在几百赫兹到几千赫兹之间。

原始数据拿回来后，第一件事不是急着做变换，而是先做几项预处理。第一步是去除直流分量，振动信号如果存在偏置，会影响小波变换后时频图的能量分布。第二步是带通滤波，滤掉明显不在轴承故障频段的干扰成分，比如低频的轴转速成分和工频干扰。第三步是尺度归一化，把信号幅值缩放到统一范围，避免不同样本的幅值差异干扰模型训练。

这些预处理步骤看起来基础，但对后续时频图质量影响很大。我曾经跳过带通滤波直接生成时频图，结果低频强干扰把故障频带的细节全掩盖了，模型怎么调都上不去。加了一阶带通滤波之后，同一套模型参数，准确率直接提升了5%。

2.2 小波基与尺度范围的确定

选择小波基函数是整个时频图生成流程中最关键的环节，也是我前期耗时最久的部分。工程上做轴承故障诊断，最常用的是Morlet小波（cmor），因为它的形态接近轴承故障产生的冲击响应，而且多了一个可调带宽参数，能够控制时频图中的频率分辨率。

实际用PyWavelets（pywt）库生成时频图时，需要设置三个核心参数：小波基名称、采样率、尺度范围。小波基名称使用cmor1.0-1.0这类格式，前面是带宽参数，后面是中心频率参数。采样率直接传入fs=12000。尺度范围则需要根据目标频率范围换算，换算公式是：实际频率等于小波中心频率乘以采样率再除以尺度值。换句话说，要覆盖0到6000Hz的频段，尺度范围就要从中心频率乘以采样率再除以6000，一路算到中心频率乘以采样率除以最低频率。

具体计算我习惯这样操作：如果小波基的中心频率是1.0，采样率12kHz，想要覆盖60Hz到6000Hz，那么最小尺度就是1.0 × 12000 ÷ 6000 = 2，最大尺度就是1.0 × 12000 ÷ 60 = 200。这个范围就可以覆盖目标频段。实际我会稍微放宽一点范围，保证边缘频率不丢失信息，但也不要放太宽，否则生成的图像里会有大量无用频带，反而稀释了模型对特征频带的注意力。

2.3 核心代码实现：生成CWT时频图

生成时频图的代码其实不长，但细节决定成败。下面是我项目里实际使用的一套完整生成流程。

import numpy as np import pywt import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from scipy.io import loadmat def cwt_to_image(signal, fs=12000, wavelet='cmor1.0-1.0', freq_min=60, freq_max=6000, image_size=224): # 尺度范围换算 fc = 1.0 # cmor1.0-1.0 的中心频率 scale_min = fc * fs / freq_max scale_max = fc * fs / freq_min scales = np.linspace(scale_min, scale_max, 128) # 执行连续小波变换 coefficients, frequencies = pywt.cwt(signal, scales, wavelet, sampling_period=1.0/fs) # 取幅值并做对数压缩，增强弱故障特征 coef_abs = np.abs(coefficients) coef_log = np.log1p(coef_abs) # 归一化到0-255范围 coef_norm = (coef_log - coef_log.min()) / (coef_log.max() - coef_log.min()) coef_uint8 = (coef_norm * 255).astype(np.uint8) # 使用伪彩色映射，并将图像resize到统一尺寸 fig = plt.figure(figsize=(4, 4), dpi=100) plt.axis('off') plt.imshow(coef_uint8, aspect='auto', cmap='jet') plt.savefig('temp_cwt.png', bbox_inches='tight', pad_inches=0) plt.close(fig) img = plt.imread('temp_cwt.png') # img 后续用 PIL 或 cv2 转为 RGB 并 resize 到 224x224 return img

有几个细节需要强调。第一，pywt.cwt返回的frequencies数组可以用来验证尺度范围换算是否正确。第二，直接用原始幅值存图，故障特征可能因为幅值差异太大而不明显，做一次对数压缩能显著增强早期微弱故障的可视化效果。第三，伪彩色映射推荐用jet，实际实验下来比gray在Transformer模型上的表现更稳定，因为色彩对比度强化了不同频带间的分辨性。

整个数据集生成时，我会把每段信号切分成固定长度的小片段，每段长度建议在4096个采样点左右。采样率12kHz时，对应约0.34秒的信号长度，这个窗口长度既包含了足够多的轴承回转周期，又不会让时频图的时间轴过于冗长。

2.4 生成时频图时容易踩的隐形坑

第一个坑是时间轴错位。切分信号时如果使用了非重叠滑窗，某些样本可能恰好把一次冲击事件切成两半，导致时频图失真。我的做法是采用50%重叠滑窗切分，每个样本保留足够的上下文信息，同时也让数据量成倍增加，对后面提精度有帮助。

第二个坑是图像resize造成的频率轴扭曲。CWT生成的时频图通常是窄长形状，如果直接粗暴压缩成224×224，高频细节会被压扁。我建议生成时就用方形的figure尺寸，让频率轴和时间轴的视觉比例接近1:1，避免resize带来的信息损失。

第三个坑是图片缓存文件管理。上面代码里用了临时文件形式，处理大批量数据时建议直接用plt.imshow配合canvas接口或者用cv2直接绘制，不要频繁读写磁盘，否则10000张图生成下来，光IO时间就够等的。我后来重构成内存中直接完成图到数组的转换，速度提升了3倍以上。

3. SwinTransformer模型落地：从配置到Pytorch实现

3.1 模型选型：用timm还是手写

SwinTransformer的实现方案，主流有两条路：一条是用timm库直接调用，一行代码加载预训练模型；另一条是从零手写Transformer模块。我的建议是，除非你要改模型结构做学术创新，否则直接用timm最省事，原因很实际——手写的SwinTransformer代码动辄几百行，里面窗口分区、相对位置编码、Shifted Window这些操作的实现细节非常多，稍不留神就会在某个张量维度上出bug，而且极难排查。

实用主义永远优先。以下是使用timm加载模型的完整代码：

import timm import torch import torch.nn as nn def build_swin_model(num_classes=4, pretrained=False, img_size=224): model = timm.create_model( 'swin_tiny_patch4_window7_224', pretrained=pretrained, num_classes=num_classes, img_size=img_size ) return model

swin_tiny_patch4_window7_224是Swin官方系列里的小型号，patch大小为4，窗口大小为7，输入图像224×224。这个配置在计算量和表达力之间比较均衡。如果你的数据集大、机器配置高，可以换成swin_base，如果追求极致轻量化，可以换成swin_small或更小的变体。

需要注意一个容易出错的地方：pretrained这个参数。时频图和ImageNet的自然图像相差很大，加载ImageNet预训练权重到底有没有帮助，我在项目中做了对比实验。结论是：预训练权重帮助不大，甚至早期会拖慢收敛，因为浅层卷积核学到的纹理特征和时频图的频率条纹结构完全不同。最终我选择了pretrained=False，从头训练反而干净利落。

3.2 核心配置参数怎么调

SwinTransformer相比CNN，超参数更多，也更敏感。我把自己实测过的一组稳定配置列出来，供参考：

drop_path_rate是我觉得最值得关注的参数。SwinTransformer在小数据集上非常容易过拟合，因为模型的参数量其实并不小。DropPath相当于正则化手段，随机丢弃部分子层路径，让模型更强壮。我在四分类任务里设置0.15，训练集和验证集的gap控制得比较好。如果你发现训练精度很高、验证精度上不去，优先怀疑DropPath太小，其次是数据增强不够。

3.3 训练循环与优化器配置

模型结构确定后，训练代码就是一个标准Pytorch流水线。但这里有几个细节和CNN训练不太一样，单独拿出来说。

优化器方面，SwinTransformer用AdamW是标配，学习率建议从1e-4起步。和CNN常用的SGD相比，AdamW对学习率没有那么敏感，但权重衰减的设置在Transformer里格外重要。weight_decay我设成0.05，太小了正则化不足，太大了模型欠拟合。

学习率调整方面，我用的是余弦退火调度器，配合前5个epoch的warmup。Transformer训练普遍需要warmup阶段，因为早期的梯度方差比较大，如果直接上大学习率，很容易掉进损失震荡的坑里。

训练精度方面，直接上自动混合精度（AMP）。SwinTransformer的显存占用在同参数量级别下比ResNet高出一大截，开启AMP之后，batch size可以翻倍，训练时间也明显缩短。以下是训练循环的简化示例：

from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50) for epoch in range(epochs): model.train() for images, labels in train_loader: images = images.cuda() labels = labels.cuda() optimizer.zero_grad() with autocast(): outputs = model(images) loss = nn.CrossEntropyLoss()(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() scheduler.step()

这段代码看着简单，却是最容易出错的地方。很多人第一次用AMP时，忘了加GradScaler，导致loss不下降或者直接NaN。另外，AMP下如果你的模型里混了自定义层且没有适配float16精度，会出现cuda error，这点在加载SwinTransformer权重时尤其要留意。

3.4 显存占用与推理加速

SwinTransformer的显存占用是真高，实话说，默认配置下224×224输入、batch size为32，一张24GB显卡才能比较从容地跑。如果显存紧张，我给出几个降显存的思路：把batch size降到16或者8；输入尺寸降到192×192；关闭AMP以外的梯度裁剪。

推理阶段的加速，我做了几项优化。模型导出时用torch.jit.trace转成TorchScript格式，推理速度能提升20%左右。更激进的做法是转成ONNX后用TensorRT加速引擎跑，但SwinTransformer的窗口注意力机制在TensorRT上支持得不算完美，有遇到算子不支持的问题，需要反复适配。我的建议是，优先用TorchScript，稳定性和兼容性最好。

4. 训练策略与实验设计：让模型真正学会"看图说话"

4.1 数据划分与防泄漏策略

轴承故障数据集有一个容易被忽略的问题：同一个工况下采集的连续信号，相邻片段之间存在极强的相关性。如果划分训练集和测试集时不做处理，只是简单随机抽样，那么训练集和测试集中很可能出现来自同一段连续信号的片段，模型学会了"记片段"而不是"学特征"，测试精度会虚高。

正确的做法是按照"运行状态记录"来划分，而不是按"片段"来划分。先把原始连续采集信号根据工况和故障状态分成若干个独立记录单元，比如按负载、转速、故障类型的不同组合来划分，然后将记录单元整体分配到训练集和测试集。我在项目中特意做了对比实验：训练测试来自同一段信号时精度达到98%，但换成了全新的运行记录时精度掉到92%，差距一目了然。

具体划分比例，我采用6:2:2——训练集60%、验证集20%、测试集20%，同时保证每个类别在各个集合中的分布基本一致。

4.2 损失函数与标签平滑

轴承故障诊断本质上是一个图像分类任务，最常规的损失函数就是交叉熵。但实际训练中我发现，直接使用标准交叉熵会导致模型过度自信，尤其是在故障程度不同的样本混杂在一起时，训练集上的loss下降得极快，验证集的表现却跟不上。

我的做法是给标签加平滑。标签平滑的思想是不再把正确类别的概率当成1，而是留出一小部分概率分配给其他类别。具体实现时，在CrossEntropyLoss中设置label_smoothing=0.1即可。这个操作非常小，却能让模型的泛化能力有一个可感知的提升，实测验证集准确率提高约1个百分点。

如果你的数据集类别不平衡严重，可以考虑在损失函数中加入类别权重，或者使用Focal Loss。Focal Loss能自动降低简单样本的权重，让模型更关注难分类的样本。但对于四分类且类别分布均衡的轴承数据，标准交叉熵加标签平滑就已经够用，不要为了炫技盲目引入复杂损失函数。

4.3 数据增强的"可以做"与"不能做"

时频图虽然不是普通照片，但它终究是一张二维图像，数据增强能帮助模型更好地泛化。我在项目中验证过几类增强手段，效果差异明显。

可以做的增强包括：随机水平裁剪和缩放，即使用RandomResizedCrop，让模型对不同时间窗口的长度有适应性；轻微的ColorJitter，调节亮度、对比度，模拟不同采集条件下幅值归一化差异带来的图像变化；随机擦除，把时频图上一小部分区域置零，模拟信号局部受干扰的情况。

不能做的增强包括：水平翻转和垂直翻转。水平翻转相当于把时间轴反向，轴承故障信号的方向性是物理意义的一部分，翻转之后物理意义就矛盾了；垂直翻转相当于把频率轴颠倒，低频变高频，这是完全错误的信号处理逻辑。在项目早期我不小心加了随机翻转，模型训练过程一直有震荡，后来排查半天才发现是这个原因。

4.4 评估指标与故障特征可视化

诊断模型的评估不能只看最终准确率，还需要看每一类故障的识别情况。轴承不同故障类型之间差异比较明显，但外圈故障和滚动体故障在某些工况下特征会部分重叠，容易混淆。我习惯用混淆矩阵来看每一个类的精确率和召回率，而不只是取一个平均指标。

做一个简单的t-SNE图也能直观呈现模型特征分布。具体做法是取SwinTransformer最后一层Transformer块输出的特征向量，经过全局平均池化后输入t-SNE降维，然后按标签着色。如果不同类别在二维空间中明显成簇，说明模型学到了具有区分度的特征；如果各类别糊成一团，那就要反思训练过程是否有问题。

5. 实战踩坑记录：文档里不会明说的细节

5.1 小波尺度差一档，精度天差地别

这是我整个项目里踩过最深的坑。一开始我偷懒，没有根据故障特征频率精确计算尺度范围，而是直接跟着网上的一个例子用了固定的尺度数组np.arange(1, 128)。这个范围在采样率12kHz、中心频率1.0的情况下，实际覆盖的频率范围是从约94Hz到12000Hz，看起来覆盖面挺广，但问题在于低频段被过度压缩了。

轴承特征频率分布在几百赫兹范围内，当我把尺度范围的上限从128改到200之后，低频段的频率分辨率瞬间提升，时频图上原本模糊成一片的故障特征带变成了清晰的条纹。同一套SwinTransformer模型，验证集准确率从88%涨到了94%。这个教训说明一个道理：时频图的质量决定了模型精度的天花板，网络再强也弥补不了输入信息的缺失。

5.2 图像缩放导致特征"压扁"

CWT生成的原始时频图通常是高大于宽的长方形，如果直接把高度从300像素压缩到224像素，频率轴的信息会被压缩变相丢失。我在测试中对比了两种方案：一种直接用transforms.Resize((224, 224))硬压缩，另一种生成图像时就把频率轴设置的宽一些、时间轴短一些，再做方形裁剪，最终准确率差了将近3个百分点。解决方案很简单：设置图像生成时的时间窗长度和频率范围，让生成的图本身就接近1:1的比例，而不是依赖后处理强行变形。

5.3 数据泄漏带来的"虚假高分"

这一点前面提过，但值得再强调一次。很多做科研的同行在出结果时发现准确率逼近100%，高兴得很，结果换一个工况的数据测试就崩盘。原因几乎都是数据划分时泄漏了。避免泄漏的方法就一条：划分数据时以连续采集记录为最小单位，而不是以单个信号片段为单位。我因为这个问题多花了整整两周的时间去排查，最后把划分逻辑重构之后才得到真正可信的结果。

5.4 AMP的坑：loss突然变NaN

在开启混合精度训练后，有一次训练过程中loss突然变成NaN，检查了好久才发现是学习率太高，导致梯度在float16精度下溢出了。解决方案是把初始学习率从5e-4降到1e-4，同时给梯度加了一个裁剪：torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)。这之后训练再没有出现过类似问题。如果你也在用AMP，建议从一开始就把这两个配置加好。

5.5 训练过程中的"memory leak"问题

还有一个不太容易想到的问题：SwinTransformer的timm实现默认会在每个epoch后缓存一部分中间变量，如果数据Loader的num_workers设置得过高，内存会持续增长。我一开始把num_workers设成16，训练10个epoch后内存直接爆掉。后来调成8，并在每个epoch结束后调用torch.cuda.empty_cache()，问题才解决。这种问题在线上的稳定训练环境里非常致命，建议尽早发现。

6. 实验结果与后续扩展思路

6.1 一套完整的实验结果对比

我们用自己采集整理的轴承数据验证了整套方案，这里给出一组代表性结果。数据集包含四类状态：正常、内圈故障、外圈故障、滚动体故障，每类400个时频图样本，共1600张图，按记录单元划分为训练集和测试集。

从结果可以看出两个要点：一是同样的模型下，CWT时频图比原始波形输入带来的提升是决定性的，这验证了前面"输入质量决定精度天花板"的判断；二是在CWT时频图输入的前提下，SwinTransformer相比ResNet50仍然有明显优势，提升了将近4个百分点。Swin-S比Swin-T提升有限，但参数量几乎翻倍，所以实际项目里我选了Swin-T，综合性价比更好。

6.2 后续可以扩展的方向

这套方案目前只是解决了"同工况下的故障分类"问题，实际工业场景中还有几个值得深挖的方向。

第一个是跨工况迁移。不同转速、不同负载下采集的振动信号，其时频图分布会整体偏移，模型泛化能力会下降。可以考虑引入域适应方法，让模型学到与工况无关的故障特征。这方面目前比较热门的做法是在SwinTransformer的特征提取层后面加一个域判别器，通过对抗训练对齐不同工况的特征分布。

第二个是连续小波尺度参数的自动搜索。目前尺度范围是手工设定的，依赖人工经验。下一步可以考虑把尺度范围当作可学习参数，用梯度下降来优化，或者用贝叶斯搜索来找到最优尺度区间。这个方法在理论上是可行的，因为CWT对尺度的微分是可以推导出来的。

第三个是模型轻量化。SwinTransformer虽然精度高，但推理速度还是偏慢，部署到边缘设备上会有压力。可以尝试知识蒸馏，用Swin-T作为教师模型，蒸馏一个轻量级的CNN学生模型，或者直接尝试SwinTransformer-Lite这类变体结构。我做过几个初版实验，蒸馏后的小模型能保留教师模型80%以上的精度提升收益，推理速度快了近3倍。

6.3 一点个人体会

做这个项目的最大收获，不是SwinTransformer本身有多强，而是认识到"数据特征表达"和"模型容量"必须匹配。很多故障诊断项目直接把原始波形塞进CNN，或者不做参数分析就直接套用ImageNet的预训练模型，最后效果不佳，却把自己的精力全都消耗在调网络结构上，步子迈得完全反了。小波时频图把一维振动信号转换成了模型更容易理解的二维图像形态，这是一个信息表达层面的优化，而SwinTransformer在此基础上把二维图像的深层模式挖掘了出来。两者结合，才让诊断系统在准确率和泛化能力上真正上一个台阶。

最后再分享一个小技巧，如果你的后续项目里也遇到了"训练集很好、测试集很差"的情况，先别急着换模型，回头看看时频图的参数设置是不是合理，很多时候问题就藏在那几个不起眼的数字里。