做轴承故障诊断这个方向有一段时间了踩了不少坑也积累了一些经验。最近把整套流程重新梳理了一遍从原始振动信号到最后模型输出决定用一篇文章把小波时频图 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清晰得多故障冲击的频带特征一眼就能分辨。这里给个表格直观对比对比维度短时傅里叶STFT连续小波变换CWT时频分辨率固定窗长分辨率恒定自适应低频高频分辨率不同对瞬态冲击信号窗内非平稳假设难满足小波基可匹配冲击形态计算复杂度较低中等偏上参数调节难点窗函数、窗长小波基、尺度范围图像特征清晰度存在模糊效应故障特征带相对清晰1.3 SwinTransformer 比 CNN 强在哪输入图像确定之后接下来就是选模型。近两年做故障诊断的主流方案依然以ResNet、VGG这类CNN为主但我在实践中的感受是CNN在时频图分类上存在两个天然短板。一个是感受野受限。CNN靠卷积核堆叠来扩大感受野但每层的感受野始终是局部的对图像中远距离依赖关系的建模能力不足。时频图上的故障特征往往不是孤立的一个点而是横跨一定时间宽度和频率宽度的条带状区域并且不同频带之间的关联可能相隔较远CNN需要非常深的网络才能捕捉到这种长程关系。另一个是层次化特征的自适应能力。CNN的池化操作是固定的下采样方式不论图像内容如何都机械地降低分辨率。SwinTransformer 引入了移动窗口注意力机制在窗口内部计算自注意力不同层之间通过窗口移位实现跨窗口信息交互。这样既保留了Transformer全局建模的能力又通过窗口化大幅降低了计算复杂度让模型的注意力能够自适应地聚焦在关键故障区域。我实测下来同一个数据集上SwinTransformerSwin-Tiny级别的准确率比ResNet50高出2到3个百分点而且在小样本场景下优势更明显。它也不是完全没有代价——训练时间确实比CNN长显存占用也更高但为了精度这些成本值得。2. 数据流水线从原始振动波形到可学习的时频图2.1 数据来源与信号预处理我们这次使用的是公开的滚动轴承故障数据集包含正常状态、内圈故障、外圈故障、滚动体故障四类样本每个类别又有不同损伤程度。数据采集时的采样率是12kHz转频范围在几十到几百赫兹故障特征频率集中在几百赫兹到几千赫兹之间。原始数据拿回来后第一件事不是急着做变换而是先做几项预处理。第一步是去除直流分量振动信号如果存在偏置会影响小波变换后时频图的能量分布。第二步是带通滤波滤掉明显不在轴承故障频段的干扰成分比如低频的轴转速成分和工频干扰。第三步是尺度归一化把信号幅值缩放到统一范围避免不同样本的幅值差异干扰模型训练。这些预处理步骤看起来基础但对后续时频图质量影响很大。我曾经跳过带通滤波直接生成时频图结果低频强干扰把故障频带的细节全掩盖了模型怎么调都上不去。加了一阶带通滤波之后同一套模型参数准确率直接提升了5%。2.2 小波基与尺度范围的确定选择小波基函数是整个时频图生成流程中最关键的环节也是我前期耗时最久的部分。工程上做轴承故障诊断最常用的是Morlet小波cmor因为它的形态接近轴承故障产生的冲击响应而且多了一个可调带宽参数能够控制时频图中的频率分辨率。实际用PyWaveletspywt库生成时频图时需要设置三个核心参数小波基名称、采样率、尺度范围。小波基名称使用cmor1.0-1.0这类格式前面是带宽参数后面是中心频率参数。采样率直接传入fs12000。尺度范围则需要根据目标频率范围换算换算公式是实际频率等于小波中心频率乘以采样率再除以尺度值。换句话说要覆盖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, fs12000, waveletcmor1.0-1.0, freq_min60, freq_max6000, image_size224): # 尺度范围换算 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_period1.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), dpi100) plt.axis(off) plt.imshow(coef_uint8, aspectauto, cmapjet) plt.savefig(temp_cwt.png, bbox_inchestight, pad_inches0) 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_classes4, pretrainedFalse, img_size224): model timm.create_model( swin_tiny_patch4_window7_224, pretrainedpretrained, num_classesnum_classes, img_sizeimg_size ) return modelswin_tiny_patch4_window7_224是Swin官方系列里的小型号patch大小为4窗口大小为7输入图像224×224。这个配置在计算量和表达力之间比较均衡。如果你的数据集大、机器配置高可以换成swin_base如果追求极致轻量化可以换成swin_small或更小的变体。需要注意一个容易出错的地方pretrained这个参数。时频图和ImageNet的自然图像相差很大加载ImageNet预训练权重到底有没有帮助我在项目中做了对比实验。结论是预训练权重帮助不大甚至早期会拖慢收敛因为浅层卷积核学到的纹理特征和时频图的频率条纹结构完全不同。最终我选择了pretrainedFalse从头训练反而干净利落。3.2 核心配置参数怎么调SwinTransformer相比CNN超参数更多也更敏感。我把自己实测过的一组稳定配置列出来供参考参数项建议值说明patch size44×4像素合并为一个token时频图细节足够embed_dim96Swin-T默认值别轻易调大window size7224×224下窗口大小为7计算量适中depths(2, 2, 6, 2)四个stage的Transformer块数量num_heads(3, 6, 12, 24)每个stage的注意力头数drop_path_rate0.1~0.2路径丢弃率防过拟合关键输入尺寸224×224如果机器允许可试256×256分类头根据类别数修改num_classes必须对应标签数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(), lr1e-4, weight_decay0.05) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50) 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_smoothing0.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_norm1.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张图按记录单元划分为训练集和测试集。模型输入特征参数量准确率测试耗时ResNet50原始波形25.6M89.2%12msResNet50CWT时频图25.6M92.4%14msSwin-TCWT时频图28.3M96.1%28msSwin-SCWT时频图49.6M96.8%35ms从结果可以看出两个要点一是同样的模型下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在此基础上把二维图像的深层模式挖掘了出来。两者结合才让诊断系统在准确率和泛化能力上真正上一个台阶。最后再分享一个小技巧如果你的后续项目里也遇到了训练集很好、测试集很差的情况先别急着换模型回头看看时频图的参数设置是不是合理很多时候问题就藏在那几个不起眼的数字里。