简介这是一份面向毕业设计、课程设计与期末大作业场景的Python轴承故障诊断项目源码核心思路为融合小波时频图与Swin Transformer模型适合计算机、自动化等相关专业学生作为完整实战范例参考。资源包共2000个文件除7个Python源程序外还包含1898张JPG时频图、40个npy与40个mat格式数据文件以及yaml、json、xml等配置和标注文件、txt与md说明文档整体大小约109.63MB目录结构清晰便于对照运行与二次开发。项目经导师指导并获高分评价源码已本地编译调试确保可运行完整覆盖振动信号处理、小波时频图生成、数据集划分、Swin Transformer网络构建训练与评估等环节有助于理解轴承故障诊断从数据到模型落地的全流程。目前已有145人浏览学习适合需要快速上手该类研究的读者。1. 小波时频图与Swin Transformer做轴承故障诊断为什么这组搭配值得复现滚动轴承故障诊断早期靠老师傅拿听诊器后来靠FFT谱找特征频率再往后是人工提取时域频域特征喂给SVM、随机森林。这套老路线的痛点是特征工程太依赖经验换一台设备、换个转速特征就不那么灵了。而“Python基于小波时频图与Swin Transformer的轴承故障诊断方法”把问题换了个思路先把一维振动信号通过连续小波变换变成二维时频图再用Swin Transformer这个视觉Transformer去自动学图里的故障模式。对正在做课程设计、毕业设计或者刚接触故障诊断方向想快速跑通一套深度学习方案的Python工程师来说这是一条有源码可拆、有公开数据集可验证、有物理含义可讲的落地路径。2. 为什么是时频图和Swin一维信号的信息缺口与Transformer的建模偏好传统方法不差但有两个瓶颈绕不开一是振动信号非平稳单靠FFT看全局频谱会丢掉故障冲击发生的时刻二是人工特征和故障类型之间的映射关系脆弱。小波时频图和Swin Transformer组合起来恰好把这两个瓶颈各自解决了一半。2.1 一维振动信号缺的是什么从傅里叶到小波时频图傅里叶变换把信号分解成不同频率的正弦波得到频谱但频率成分随时间怎么变化是看不到的。短时傅里叶变换STFT加了窗能看时变特性可窗口一固定低频和高频的分辨率就互相牵制——窗口短时间分辨率高但频率分辨率差窗口长反过来。轴承故障信号里的冲击成分恰恰是宽带高频、持续时间短的瞬态用STFT容易把冲击细节抹平。连续小波变换CWT的思路不一样用一个小波母函数做尺度伸缩和平移尺度小对应高频、时间窗窄尺度大对应低频、时间窗宽天然就是多分辨率分析。对一段振动信号做CWT得到的是一个“时间-尺度频率”的二维系数矩阵幅度反映该时刻该频带的能量强弱。把系数矩阵映射成图像像素值就是小波时频图。早期微弱故障的特征往往就在某个共振频带附近以边带形式出现时频图能把这种“什么时候、在哪个频带、能量多强”的信息同时呈现出来。比起直接拿原始波形或FFT谱当输入时频图保留了时变信息这是它最值钱的地方。常见做法是用Python里的PyWavelets库做CWT尺度序列用对数分布覆盖关注的频带然后取系数的模做归一化保存成图片。代码量不大但参数选错时频图就是一团糊这个后面单独讲。2.2 为什么选Swin Transformer而不是CNN或ResNet滚动轴承的时频图有个特点故障特征频率和它的倍频、边带在图上表现为多个离散的水平亮带彼此间隔有物理规律。要识别这种模式模型需要同时看到“局部频带里有没有共振能量”和“这些亮带之间的间隔关系”。ResNet这类CNN擅长提局部纹理但感受野受卷积核大小和层数限制对长距离位置关系的建模偏弱。Vision TransformerViT用全局自注意力能看全图但patch数量一多计算量随序列长度平方增长训练和推理都不好伺候。Swin Transformer的折中方案是分层窗口自注意力在局部窗口内做自注意力W-MSA窗口之间用移位窗口Shifted Window来传递信息再配合Patch Merging逐层下采样形成类似CNN的多尺度金字塔。这样既保留了Transformer的长程建模能力又把计算复杂度控制在和图像分辨率线性相关所以Swin从2021年提出后就一直是视觉任务里性价比很高的backbone。对故障诊断场景来说用Swin并不是什么玄学它就是让模型能自由建立“图上任意两个能量带之间的关联”而这种关联恰好在故障模式里是强物理特征比如内圈故障时转频调制产生的边带间隔。可以说Swin把“找特征”这件事的搜索空间放大了一截尤其在样本量还可以、图像尺寸不超过224或256的时频图数据集上它比同等规模的ResNet更容易吃满信息。2.3 整体流程与数据划分先想清楚再写代码整套方案可以拆成五个环节振动信号读取、滑动窗口切段、CWT时频图生成、按工况划分数据集、Swin Transformer训练评估。实际编码顺序也是这个顺序每步都有独立产出。# 划分数据集的核心原则按“数据段”划分不按“单窗口”划分 def split_by_segment(files, train_ratio0.7): # files: 每个元素是一个振动信号文件的路径对应一个连续采集的数据段 # 同一个文件切出来的所有窗口必须全部进入训练集或测试集 from sklearn.model_selection import train_test_split train_files, test_files train_test_split( files, train_sizetrain_ratio, stratify[get_label(f) for f in files], # 按故障类别分层 random_state42 ) return train_files, test_files这里的逻辑说明一下很多人拿到数据后直接把所有窗口图混合起来随机划分验证集准确率能到99%一上真实数据就崩。原因是相邻窗口来自同一段信号图像几乎一样模型等于提前见过答案。正确做法是先按采集文件数据段划分再在每个文件内部滑窗生成图片。stratify参数按故障类别分层保证每一类在训练和测试里都有代表。代码里get_label需要你自己根据文件命名规则实现一般是正则表达式提取故障类型字符串。数据量方面每类故障建议不少于800到1000张时频图否则Swin这种参数量的模型容易过拟合。公开的CWRU、东南大学轴承数据集、XJTU-SY轴承数据都够用如果自己采数据采样率至少设到12kHz以上并且记录转速后面算特征频率要用。3. 从振动信号到小波时频图CWT参数与批量生成代码时频图的质量直接决定分类上限。这一章给出可直接抄的CWT生成代码以及窗口长度、重叠率、尺度序列这些参数的选取方法。3.1 用PyWavelets生成CWT时频图最小可运行代码生成一张时频图核心就是pywt.cwt一行但前后处理决定了这张图能不能用。import numpy as np import pywt import matplotlib.pyplot as plt def cwt_to_image(signal, fs, waveletcmor, n_scales64, target_size(224, 224)): # 尺度序列对数均匀分布覆盖约1Hz到Nyquist频率 scales np.logspace(np.log10(1), np.log10(fs / 2), n_scales) # 连续小波变换返回 (n_scales, n_times) 的复系数矩阵 coeffs, _ pywt.cwt(signal, scales, wavelet, sampling_period1.0 / fs) # 取模并做对数压缩避免低频系数淹没高频细节 log_abs np.log1p(np.abs(coeffs)) # 归一化到 [0, 255]作为图像保存 norm (log_abs - log_abs.min()) / (log_abs.max() - log_abs.min() 1e-8) img (norm * 255).astype(np.uint8) # 缩放到统一尺寸Swin的patch embedding需要固定分辨率 from PIL import Image img Image.fromarray(img).resize(target_size, Image.Resampling.LANCZOS) return np.asarray(img)逐行说明scales用logspace生成从1到fs/2覆盖从低频冲击包络到高频共振带64个尺度对224x224的图够用想要更细腻可以提到128。wavelet选cmor复morlet小波它对振荡信号的分析精度好相频信息也能保留想用实数小波就选mexh但对边带类特征识别稍弱。np.log1p是这代码里最关键的一步——CWT系数动态范围极大低频主轴能量可能是高频细节的上万倍直接min-max归一化会把微小故障特征压成全黑或全灰取对数后能量对比关系才合理。缩放到224x224和Swin-T的输入尺寸对齐避免patch切分不均。3.2 滑动窗口切段长度、重叠率与样本量怎么配生成时频图之前要先给连续信号切段。窗口长度拍脑袋定会出问题太短看不全一个完整的故障冲击周期太长则一段样本里包含的工况变化太复杂类别特征被平均掉。def sliding_window(signal, win_len, overlap0.5): stride int(win_len * (1 - overlap)) if stride 1: raise ValueError(overlap过高stride不能小于1) n_windows (len(signal) - win_len) // stride 1 return [signal[i * stride : i * stride win_len] for i in range(n_windows)] # 参数估算已知转频 fr希望每个窗口覆盖至少30圈 # fs 12000, fr 30Hz - win_len int(fs * 30 / fr) 12000 win_len int(fs * 30 / fr) windows sliding_window(raw_signal, win_lenwin_len, overlap0.5)窗口长度按转频来估算是最稳的每个窗口至少包含20到30个旋转周期的振动数据这样故障冲击以稳定的频率出现在窗口内时频图才有规律的横向条纹。重叠率常见取50%样本量不够就提到75%但要注意重叠率越高相邻窗口图像越相似训练时间长了容易让模型对相位位置过拟合。这段代码的边界条件是n_windows必须至少为1同时重叠率不能导致stride为零实际做批量生成时建议加个断言。3.3 批量生成与目录组织直接喂给torchvision的ImageFolder训练Swin时最省事的做法是把数据集整理成ImageFolder标准结构每类故障一个子目录图片直接是jpg或png。目录组织如下目录路径内容标签data/train/Normal正常状态时频图0data/train/InnerRace内圈故障时频图1data/train/OuterRace外圈故障时频图2data/train/Roller滚动体故障时频图3data/val/Normal验证集0批量生成的代码逻辑是遍历每个信号的每个窗口调用cwt_to_image保存成文件文件名里带上信号编号和窗口序号。import os from pathlib import Path def build_dataset(signal_files, out_root, fs, win_len, overlap): # signal_files: list[dict]每个元素里有 path 和 label 两个键 # 按数据段划分后训练和验证的文件列表是分开传入的 for row in signal_files: raw np.loadtxt(row[path]) # 从csv/txt读取振动信号 seg_name Path(row[path]).stem save_dir Path(out_root) / row[label] save_dir.mkdir(parentsTrue, exist_okTrue) windows sliding_window(raw, win_len, overlap) for idx, seg in enumerate(windows): img cwt_to_image(seg, fs) save_path save_dir / f{seg_name}_w{idx:04d}.png Image.fromarray(img).save(save_path)这段代码里注意两点。第一out_root要分成train和val两个根目录而且传入的signal_files必须已经是按数据段划分好的列表不能把同一个信号文件的窗口同时散落到两个集合。第二文件名用seg_name加_w编号是为了万一后面要回溯某个测试样本来自哪一段信号时能找回来这是个排查问题的好习惯。图片格式统一png无损且尺寸可控。4. Swin Transformer训练源码从模型加载到损失函数三个关键改动有了数据集接下来就是把Swin跑起来。这里不推荐手写公式复现Swin结构直接用timm库加载官方权重把分类头换掉然后调训练策略。4.1 用timm加载Swin-T预训练权重怎么接上Swin Transformer的PyTorch官方实现存在SwinTransformer仓库里结构完整但使用起来有些笨重。我一般直接用timm的封装一行就能建好模型。import timm import torch.nn as nn model timm.create_model( swin_tiny_patch4_window7_224, pretrainedTrue, num_classes4 ) # 查看分类头结构 in_features model.head.in_features model.head nn.Linear(in_features, 4)参数说明模型名字swin_tiny_patch4_window7_224拆开看patch4表示每个patch是4x4像素window7表示注意力窗口是7x7224是输入分辨率。这套配置是Swin-T的默认标准预训练权重是在ImageNet-22K上训过的对自然图像的颜色、纹理、边缘有很强的先验。时频图虽然不是自然图像但它也是“二维纹理局部结构”的模式前几层学的边缘和纹理滤波器可以直接迁移所以用预训练权重比从零训练收敛快得多尤其当你的故障数据集只有几千张图时。这里有个坑num_classes4直接传进去后timm会自动帮你换掉最后的分类头但如果你先创建了num_classes1000的模型再手动改head要注意model.head的名字是head而不是fc改错属性名的话加载预训练权重时后面会报参数不匹配。4.2 训练主循环数据集、优化器、学习率策略一次配齐训练脚本的结构是标准的PyTorch流程工程上真正影响结果的是两个细节一是Swin这类Transformer对学习率比CNN敏感二是AdamW加余弦退火是它最稳的组合。import torch from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder from torchvision import transforms import timm, torch.nn as nn transform transforms.Compose([ transforms.RandomHorizontalFlip(p0.5), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) train_ds ImageFolder(data/train, transformtransform) val_ds ImageFolder(data/val, transformtransform) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue) model timm.create_model(swin_tiny_patch4_window7_224, pretrainedTrue, num_classes4) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay0.05) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max30, eta_min1e-6) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) for epoch in range(30): model.train() running_loss 0.0 for x, y in train_loader: x, y x.to(device), y.to(device) optimizer.zero_grad() logits model(x) loss criterion(logits, y) loss.backward() optimizer.step() running_loss loss.item() scheduler.step() # 每个epoch结束后在验证集上评估一次并保存best print(fEpoch {epoch1}, loss: {running_loss / len(train_loader):.4f})几个参数值得说道。学习率1e-4是Swin微调的常见起点比ResNet常用的1e-3低一个量级因为Transformer的自注意力对参数扰动敏感lr给大了前几轮就会震荡。weight_decay设到0.05这比CNN常用的1e-4大得多是AdamW针对Transformer调出来的经验值起到抑制高频噪声特征的作用。CosineAnnealingLR把学习率从1e-4平滑降到1e-6相比StepLR阶跃下降余弦曲线在训练后期更稳最终精度一般高0.5到1个百分点。batch_size根据显存定显存不够时优先减到16或8而不是硬撑到OOM。4.3 类别不平衡与损失函数加权CrossEntropy比调数据快轴承故障数据里最常见的不平衡是正常样本远多于故障样本或者某种故障类型采集时间短、样本量只有其他类的三分之一。这会让模型把样本多的类学得更好样本少的类准确率塌掉。from sklearn.utils.class_weight import compute_class_weight import numpy as np # train_ds.targets 是 ImageFolder 自动生成的标签列表 class_weights compute_class_weight( balanced, classesnp.unique(train_ds.targets), ytrain_ds.targets ) class_weights torch.tensor(class_weights, dtypetorch.float32).to(device) criterion nn.CrossEntropyLoss(weightclass_weights)compute_class_weight会按“样本数的倒数比例”计算每个类别的权重样本少的故障类权重自动变大损失函数对它的错判给更高的惩罚。这个改动的性价比很高只改一行初始化代码不需要动数据集也不需要在DataLoader里写sampler做采样。如果类别不平衡特别严重比如正常样本占90%可以在此基础上叠加LabelSmoothing把CrossEntropyLoss的label_smoothing0.1打开防止模型对训练集标签过于自信提升一点泛化。注意加了class_weight后验证集上的准确率不能直接横向对比要配合后面的混淆矩阵看每一类的召回率。5. 避坑排查5个让训练翻车的细节与修复这套方案我在不同数据集上跑过好几轮翻车的点高度集中在这5个地方。每一条都是“现象→原因→解决”的结构踩过一次就能避开。5.1 时频图全黑或全白模型怎么训都欠拟合现象生成的时频图看起来几乎是一个纯色块边界能看到几条模糊的线Swin训练十几个epoch准确率还在20%左右徘徊。原因CWT系数矩阵里低频能量极大直接线性归一化时最大值被低频占满高频故障冲击的数值相对太小被压到了像素值0附近表现在图上就是大片黑色。解决用对数压缩后再归一化也就是前面代码里的np.log1p同时把colormap固定为gray而不是jet避免颜色映射对模型引入不必要的噪音。改完后图上的横向亮带应该清晰可见低频部分不再是一坨全白。5.2 验证集99%测试集崩到70%划分泄露现象训练时验证集准确率一路涨到99%训练曲线很漂亮但把数据换成同一台设备新采集的信号准确率掉20个点以上。原因这是最典型的翻车点——滑动窗口重叠率高了之后训练集和验证集里存在来自同一段采集信号的窗口图像内容高度相似模型实际上是在“背图”。解决回到第2章的split_by_segment先按采集文件把数据分成若干个段再按数据段划分训练/验证/测试同一个段的窗口永远只能出现在一边。严格的测试集应该是在训练结束后用另一段从未参与过CWT和划分的信号文件来构建模拟真实新样本输入。5.3 加载预训练权重报RuntimeError: size mismatch现象执行model.load_state_dict(torch.load(swin_tiny.pth))抛错提示分类头head.weight尺寸不一致。原因Swin预训练权重输出是1000类ImageNet分类头而本地模型分类头是4类shape对不上有时还会顺带报norm层的参数错位。解决不要直接load整个state_dict使用timmcreate_model时就传num_classes4它内部会正确地跳过分类头、保留backbone权重。如果非要手动load用strictFalse加载并只load不含head的权重项但这样不推荐容易把后续结构改动埋下隐患。5.4 显存不够、batch设很小或者loss剧烈震荡现象2080Ti 11G显存跑batch32直接OOM调到8才能跑但loss波动很大验证集精度的方差也变大。原因batch太小导致梯度估计噪声大Swin的BN层在小batch下统计量不稳定。解决显存不够时优先用混合精度训练torch.cuda.amp常见可以把batch翻倍batch仍然小就增大梯度累积步数模拟更大batch的梯度更新。另外把num_workers调到4以上pin_memory打开能减少GPU空闲等待缓解训练的不确定性。5.5 混淆矩阵里某一类故障被模型全部认成另一类现象看混淆矩阵内圈故障那行几乎全落在外圈故障列其他类正常。原因两类故障的特征频率在尺度序列覆盖的低频区间重合或者是该故障样本量过少、时频图模式太单一模型没有见过足够多样的相位和负载变化。解决检查尺度序列是否覆盖了该类故障的特征频率区间低转速下特征频率可能低于1Hzscales下限要从1改到0.5样本少就提高窗口重叠率到75%或者做数据增强给时频图加随机噪声、时间小幅平移让模型对相位变化不敏感。改完后单独看那一类的召回率有没有上升而不是整体准确率。6. 结果验证与部署边界特征可视化、特征频率对照与模型导出模型训完不代表方案结束还要确认它是靠真正的物理特征分类而不是数据集里的偶然线索。这一步也是答辩和项目汇报时最能加分的部分。6.1 t-SNE特征可视化与混淆矩阵看模型有没有学到分簇from sklearn.manifold import TSNE import matplotlib.pyplot as plt # 把模型去掉head取倒数第二层的特征向量 model.eval() features, labels [], [] with torch.no_grad(): for x, y in val_loader: x x.to(device) feat model.forward_features(x) # timm模型内置的feature方法 features.append(feat.mean(dim(1, 2)).cpu()) labels.append(y) features torch.cat(features).numpy() tsne TSNE(n_components2, perplexity30, random_state0).fit_transform(features) plt.scatter(tsne[:, 0], tsne[:, 1], ctorch.cat(labels), cmaptab10, s10) plt.savefig(tsne_check.png)forward_features输出的是Swin最后一个stage的特征图形状是(B, H/32, W/32, C)对空间维度求均值得到一个向量代表整张图的全局特征。t-SNE图里四类颜色明显分成四簇说明模型学到的是类间可分的模式如果某一类分簇散乱优先怀疑训练数据本身就不够丰富。配合混淆矩阵看每一个类的召回率比单看准确率更能暴露问题。6.2 特征频率对照表用物理规律验证诊断可信度训练结束后要做最后一个验证把模型最关注的频带和轴承故障特征频率的理论值对照。滚动轴承故障特征频率公式如下需要知道转频fr、滚动体个数n、滚动体直径d、节径D和接触角α故障位置特征频率公式典型表现时频图外圈 BPFOn/2 * fr * (1 - d/D * cosα)固定亮带无调制内圈 BPFIn/2 * fr * (1 d/D * cosα)亮带带边带有转频调制滚动体 BSFD/(2d) * fr * (1 - (d/D)^2 * cos²α)双频亮带常有间隔不均保持架 FTFfr/2 * (1 - d/D * cosα)极低频亮带对照方法是看测试集里该类样本的时频图确认图中是否有对应特征频率及其倍频的亮带。模型如果主要靠这些带分类那它的决策有明显物理依据不是纯粹在拟合数据纹理。6.3 TorchScript导出部署前的边界锁定model.eval() dummy torch.randn(1, 3, 224, 224).to(device) traced torch.jit.trace(model, dummy) traced.save(bearing_swin.pt)导出时最容易被忽略的是预处理参数必须和训练时完全一致。比如训练时用了ImageNet的mean/std归一化部署端也要用同一组数值输入尺寸固定224x224cam或者上位机送进来的图必须先缩放再归一化。我早前踩过一次训练归一化没做、部署时做了测试集准确率从96%掉到88%排查半天发现是预处理链路不一致。这套方案真正的价值不是把准确率堆到99%而是让故障诊断有从信号到图像的完整链路、有物理公式可解释、有可落地的模型导出方案在这个方向上按数据段划分、按特征频率验证、按预处理锁定这三件事养成的习惯至今都在帮我省时间。希望帮到你。本文还有配套的精品资源点击获取
Python小波时频图与Swin Transformer轴承故障诊断实战
简介这是一份面向毕业设计、课程设计与期末大作业场景的Python轴承故障诊断项目源码核心思路为融合小波时频图与Swin Transformer模型适合计算机、自动化等相关专业学生作为完整实战范例参考。资源包共2000个文件除7个Python源程序外还包含1898张JPG时频图、40个npy与40个mat格式数据文件以及yaml、json、xml等配置和标注文件、txt与md说明文档整体大小约109.63MB目录结构清晰便于对照运行与二次开发。项目经导师指导并获高分评价源码已本地编译调试确保可运行完整覆盖振动信号处理、小波时频图生成、数据集划分、Swin Transformer网络构建训练与评估等环节有助于理解轴承故障诊断从数据到模型落地的全流程。目前已有145人浏览学习适合需要快速上手该类研究的读者。1. 小波时频图与Swin Transformer做轴承故障诊断为什么这组搭配值得复现滚动轴承故障诊断早期靠老师傅拿听诊器后来靠FFT谱找特征频率再往后是人工提取时域频域特征喂给SVM、随机森林。这套老路线的痛点是特征工程太依赖经验换一台设备、换个转速特征就不那么灵了。而“Python基于小波时频图与Swin Transformer的轴承故障诊断方法”把问题换了个思路先把一维振动信号通过连续小波变换变成二维时频图再用Swin Transformer这个视觉Transformer去自动学图里的故障模式。对正在做课程设计、毕业设计或者刚接触故障诊断方向想快速跑通一套深度学习方案的Python工程师来说这是一条有源码可拆、有公开数据集可验证、有物理含义可讲的落地路径。2. 为什么是时频图和Swin一维信号的信息缺口与Transformer的建模偏好传统方法不差但有两个瓶颈绕不开一是振动信号非平稳单靠FFT看全局频谱会丢掉故障冲击发生的时刻二是人工特征和故障类型之间的映射关系脆弱。小波时频图和Swin Transformer组合起来恰好把这两个瓶颈各自解决了一半。2.1 一维振动信号缺的是什么从傅里叶到小波时频图傅里叶变换把信号分解成不同频率的正弦波得到频谱但频率成分随时间怎么变化是看不到的。短时傅里叶变换STFT加了窗能看时变特性可窗口一固定低频和高频的分辨率就互相牵制——窗口短时间分辨率高但频率分辨率差窗口长反过来。轴承故障信号里的冲击成分恰恰是宽带高频、持续时间短的瞬态用STFT容易把冲击细节抹平。连续小波变换CWT的思路不一样用一个小波母函数做尺度伸缩和平移尺度小对应高频、时间窗窄尺度大对应低频、时间窗宽天然就是多分辨率分析。对一段振动信号做CWT得到的是一个“时间-尺度频率”的二维系数矩阵幅度反映该时刻该频带的能量强弱。把系数矩阵映射成图像像素值就是小波时频图。早期微弱故障的特征往往就在某个共振频带附近以边带形式出现时频图能把这种“什么时候、在哪个频带、能量多强”的信息同时呈现出来。比起直接拿原始波形或FFT谱当输入时频图保留了时变信息这是它最值钱的地方。常见做法是用Python里的PyWavelets库做CWT尺度序列用对数分布覆盖关注的频带然后取系数的模做归一化保存成图片。代码量不大但参数选错时频图就是一团糊这个后面单独讲。2.2 为什么选Swin Transformer而不是CNN或ResNet滚动轴承的时频图有个特点故障特征频率和它的倍频、边带在图上表现为多个离散的水平亮带彼此间隔有物理规律。要识别这种模式模型需要同时看到“局部频带里有没有共振能量”和“这些亮带之间的间隔关系”。ResNet这类CNN擅长提局部纹理但感受野受卷积核大小和层数限制对长距离位置关系的建模偏弱。Vision TransformerViT用全局自注意力能看全图但patch数量一多计算量随序列长度平方增长训练和推理都不好伺候。Swin Transformer的折中方案是分层窗口自注意力在局部窗口内做自注意力W-MSA窗口之间用移位窗口Shifted Window来传递信息再配合Patch Merging逐层下采样形成类似CNN的多尺度金字塔。这样既保留了Transformer的长程建模能力又把计算复杂度控制在和图像分辨率线性相关所以Swin从2021年提出后就一直是视觉任务里性价比很高的backbone。对故障诊断场景来说用Swin并不是什么玄学它就是让模型能自由建立“图上任意两个能量带之间的关联”而这种关联恰好在故障模式里是强物理特征比如内圈故障时转频调制产生的边带间隔。可以说Swin把“找特征”这件事的搜索空间放大了一截尤其在样本量还可以、图像尺寸不超过224或256的时频图数据集上它比同等规模的ResNet更容易吃满信息。2.3 整体流程与数据划分先想清楚再写代码整套方案可以拆成五个环节振动信号读取、滑动窗口切段、CWT时频图生成、按工况划分数据集、Swin Transformer训练评估。实际编码顺序也是这个顺序每步都有独立产出。# 划分数据集的核心原则按“数据段”划分不按“单窗口”划分 def split_by_segment(files, train_ratio0.7): # files: 每个元素是一个振动信号文件的路径对应一个连续采集的数据段 # 同一个文件切出来的所有窗口必须全部进入训练集或测试集 from sklearn.model_selection import train_test_split train_files, test_files train_test_split( files, train_sizetrain_ratio, stratify[get_label(f) for f in files], # 按故障类别分层 random_state42 ) return train_files, test_files这里的逻辑说明一下很多人拿到数据后直接把所有窗口图混合起来随机划分验证集准确率能到99%一上真实数据就崩。原因是相邻窗口来自同一段信号图像几乎一样模型等于提前见过答案。正确做法是先按采集文件数据段划分再在每个文件内部滑窗生成图片。stratify参数按故障类别分层保证每一类在训练和测试里都有代表。代码里get_label需要你自己根据文件命名规则实现一般是正则表达式提取故障类型字符串。数据量方面每类故障建议不少于800到1000张时频图否则Swin这种参数量的模型容易过拟合。公开的CWRU、东南大学轴承数据集、XJTU-SY轴承数据都够用如果自己采数据采样率至少设到12kHz以上并且记录转速后面算特征频率要用。3. 从振动信号到小波时频图CWT参数与批量生成代码时频图的质量直接决定分类上限。这一章给出可直接抄的CWT生成代码以及窗口长度、重叠率、尺度序列这些参数的选取方法。3.1 用PyWavelets生成CWT时频图最小可运行代码生成一张时频图核心就是pywt.cwt一行但前后处理决定了这张图能不能用。import numpy as np import pywt import matplotlib.pyplot as plt def cwt_to_image(signal, fs, waveletcmor, n_scales64, target_size(224, 224)): # 尺度序列对数均匀分布覆盖约1Hz到Nyquist频率 scales np.logspace(np.log10(1), np.log10(fs / 2), n_scales) # 连续小波变换返回 (n_scales, n_times) 的复系数矩阵 coeffs, _ pywt.cwt(signal, scales, wavelet, sampling_period1.0 / fs) # 取模并做对数压缩避免低频系数淹没高频细节 log_abs np.log1p(np.abs(coeffs)) # 归一化到 [0, 255]作为图像保存 norm (log_abs - log_abs.min()) / (log_abs.max() - log_abs.min() 1e-8) img (norm * 255).astype(np.uint8) # 缩放到统一尺寸Swin的patch embedding需要固定分辨率 from PIL import Image img Image.fromarray(img).resize(target_size, Image.Resampling.LANCZOS) return np.asarray(img)逐行说明scales用logspace生成从1到fs/2覆盖从低频冲击包络到高频共振带64个尺度对224x224的图够用想要更细腻可以提到128。wavelet选cmor复morlet小波它对振荡信号的分析精度好相频信息也能保留想用实数小波就选mexh但对边带类特征识别稍弱。np.log1p是这代码里最关键的一步——CWT系数动态范围极大低频主轴能量可能是高频细节的上万倍直接min-max归一化会把微小故障特征压成全黑或全灰取对数后能量对比关系才合理。缩放到224x224和Swin-T的输入尺寸对齐避免patch切分不均。3.2 滑动窗口切段长度、重叠率与样本量怎么配生成时频图之前要先给连续信号切段。窗口长度拍脑袋定会出问题太短看不全一个完整的故障冲击周期太长则一段样本里包含的工况变化太复杂类别特征被平均掉。def sliding_window(signal, win_len, overlap0.5): stride int(win_len * (1 - overlap)) if stride 1: raise ValueError(overlap过高stride不能小于1) n_windows (len(signal) - win_len) // stride 1 return [signal[i * stride : i * stride win_len] for i in range(n_windows)] # 参数估算已知转频 fr希望每个窗口覆盖至少30圈 # fs 12000, fr 30Hz - win_len int(fs * 30 / fr) 12000 win_len int(fs * 30 / fr) windows sliding_window(raw_signal, win_lenwin_len, overlap0.5)窗口长度按转频来估算是最稳的每个窗口至少包含20到30个旋转周期的振动数据这样故障冲击以稳定的频率出现在窗口内时频图才有规律的横向条纹。重叠率常见取50%样本量不够就提到75%但要注意重叠率越高相邻窗口图像越相似训练时间长了容易让模型对相位位置过拟合。这段代码的边界条件是n_windows必须至少为1同时重叠率不能导致stride为零实际做批量生成时建议加个断言。3.3 批量生成与目录组织直接喂给torchvision的ImageFolder训练Swin时最省事的做法是把数据集整理成ImageFolder标准结构每类故障一个子目录图片直接是jpg或png。目录组织如下目录路径内容标签data/train/Normal正常状态时频图0data/train/InnerRace内圈故障时频图1data/train/OuterRace外圈故障时频图2data/train/Roller滚动体故障时频图3data/val/Normal验证集0批量生成的代码逻辑是遍历每个信号的每个窗口调用cwt_to_image保存成文件文件名里带上信号编号和窗口序号。import os from pathlib import Path def build_dataset(signal_files, out_root, fs, win_len, overlap): # signal_files: list[dict]每个元素里有 path 和 label 两个键 # 按数据段划分后训练和验证的文件列表是分开传入的 for row in signal_files: raw np.loadtxt(row[path]) # 从csv/txt读取振动信号 seg_name Path(row[path]).stem save_dir Path(out_root) / row[label] save_dir.mkdir(parentsTrue, exist_okTrue) windows sliding_window(raw, win_len, overlap) for idx, seg in enumerate(windows): img cwt_to_image(seg, fs) save_path save_dir / f{seg_name}_w{idx:04d}.png Image.fromarray(img).save(save_path)这段代码里注意两点。第一out_root要分成train和val两个根目录而且传入的signal_files必须已经是按数据段划分好的列表不能把同一个信号文件的窗口同时散落到两个集合。第二文件名用seg_name加_w编号是为了万一后面要回溯某个测试样本来自哪一段信号时能找回来这是个排查问题的好习惯。图片格式统一png无损且尺寸可控。4. Swin Transformer训练源码从模型加载到损失函数三个关键改动有了数据集接下来就是把Swin跑起来。这里不推荐手写公式复现Swin结构直接用timm库加载官方权重把分类头换掉然后调训练策略。4.1 用timm加载Swin-T预训练权重怎么接上Swin Transformer的PyTorch官方实现存在SwinTransformer仓库里结构完整但使用起来有些笨重。我一般直接用timm的封装一行就能建好模型。import timm import torch.nn as nn model timm.create_model( swin_tiny_patch4_window7_224, pretrainedTrue, num_classes4 ) # 查看分类头结构 in_features model.head.in_features model.head nn.Linear(in_features, 4)参数说明模型名字swin_tiny_patch4_window7_224拆开看patch4表示每个patch是4x4像素window7表示注意力窗口是7x7224是输入分辨率。这套配置是Swin-T的默认标准预训练权重是在ImageNet-22K上训过的对自然图像的颜色、纹理、边缘有很强的先验。时频图虽然不是自然图像但它也是“二维纹理局部结构”的模式前几层学的边缘和纹理滤波器可以直接迁移所以用预训练权重比从零训练收敛快得多尤其当你的故障数据集只有几千张图时。这里有个坑num_classes4直接传进去后timm会自动帮你换掉最后的分类头但如果你先创建了num_classes1000的模型再手动改head要注意model.head的名字是head而不是fc改错属性名的话加载预训练权重时后面会报参数不匹配。4.2 训练主循环数据集、优化器、学习率策略一次配齐训练脚本的结构是标准的PyTorch流程工程上真正影响结果的是两个细节一是Swin这类Transformer对学习率比CNN敏感二是AdamW加余弦退火是它最稳的组合。import torch from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder from torchvision import transforms import timm, torch.nn as nn transform transforms.Compose([ transforms.RandomHorizontalFlip(p0.5), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) train_ds ImageFolder(data/train, transformtransform) val_ds ImageFolder(data/val, transformtransform) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue) model timm.create_model(swin_tiny_patch4_window7_224, pretrainedTrue, num_classes4) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay0.05) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max30, eta_min1e-6) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) for epoch in range(30): model.train() running_loss 0.0 for x, y in train_loader: x, y x.to(device), y.to(device) optimizer.zero_grad() logits model(x) loss criterion(logits, y) loss.backward() optimizer.step() running_loss loss.item() scheduler.step() # 每个epoch结束后在验证集上评估一次并保存best print(fEpoch {epoch1}, loss: {running_loss / len(train_loader):.4f})几个参数值得说道。学习率1e-4是Swin微调的常见起点比ResNet常用的1e-3低一个量级因为Transformer的自注意力对参数扰动敏感lr给大了前几轮就会震荡。weight_decay设到0.05这比CNN常用的1e-4大得多是AdamW针对Transformer调出来的经验值起到抑制高频噪声特征的作用。CosineAnnealingLR把学习率从1e-4平滑降到1e-6相比StepLR阶跃下降余弦曲线在训练后期更稳最终精度一般高0.5到1个百分点。batch_size根据显存定显存不够时优先减到16或8而不是硬撑到OOM。4.3 类别不平衡与损失函数加权CrossEntropy比调数据快轴承故障数据里最常见的不平衡是正常样本远多于故障样本或者某种故障类型采集时间短、样本量只有其他类的三分之一。这会让模型把样本多的类学得更好样本少的类准确率塌掉。from sklearn.utils.class_weight import compute_class_weight import numpy as np # train_ds.targets 是 ImageFolder 自动生成的标签列表 class_weights compute_class_weight( balanced, classesnp.unique(train_ds.targets), ytrain_ds.targets ) class_weights torch.tensor(class_weights, dtypetorch.float32).to(device) criterion nn.CrossEntropyLoss(weightclass_weights)compute_class_weight会按“样本数的倒数比例”计算每个类别的权重样本少的故障类权重自动变大损失函数对它的错判给更高的惩罚。这个改动的性价比很高只改一行初始化代码不需要动数据集也不需要在DataLoader里写sampler做采样。如果类别不平衡特别严重比如正常样本占90%可以在此基础上叠加LabelSmoothing把CrossEntropyLoss的label_smoothing0.1打开防止模型对训练集标签过于自信提升一点泛化。注意加了class_weight后验证集上的准确率不能直接横向对比要配合后面的混淆矩阵看每一类的召回率。5. 避坑排查5个让训练翻车的细节与修复这套方案我在不同数据集上跑过好几轮翻车的点高度集中在这5个地方。每一条都是“现象→原因→解决”的结构踩过一次就能避开。5.1 时频图全黑或全白模型怎么训都欠拟合现象生成的时频图看起来几乎是一个纯色块边界能看到几条模糊的线Swin训练十几个epoch准确率还在20%左右徘徊。原因CWT系数矩阵里低频能量极大直接线性归一化时最大值被低频占满高频故障冲击的数值相对太小被压到了像素值0附近表现在图上就是大片黑色。解决用对数压缩后再归一化也就是前面代码里的np.log1p同时把colormap固定为gray而不是jet避免颜色映射对模型引入不必要的噪音。改完后图上的横向亮带应该清晰可见低频部分不再是一坨全白。5.2 验证集99%测试集崩到70%划分泄露现象训练时验证集准确率一路涨到99%训练曲线很漂亮但把数据换成同一台设备新采集的信号准确率掉20个点以上。原因这是最典型的翻车点——滑动窗口重叠率高了之后训练集和验证集里存在来自同一段采集信号的窗口图像内容高度相似模型实际上是在“背图”。解决回到第2章的split_by_segment先按采集文件把数据分成若干个段再按数据段划分训练/验证/测试同一个段的窗口永远只能出现在一边。严格的测试集应该是在训练结束后用另一段从未参与过CWT和划分的信号文件来构建模拟真实新样本输入。5.3 加载预训练权重报RuntimeError: size mismatch现象执行model.load_state_dict(torch.load(swin_tiny.pth))抛错提示分类头head.weight尺寸不一致。原因Swin预训练权重输出是1000类ImageNet分类头而本地模型分类头是4类shape对不上有时还会顺带报norm层的参数错位。解决不要直接load整个state_dict使用timmcreate_model时就传num_classes4它内部会正确地跳过分类头、保留backbone权重。如果非要手动load用strictFalse加载并只load不含head的权重项但这样不推荐容易把后续结构改动埋下隐患。5.4 显存不够、batch设很小或者loss剧烈震荡现象2080Ti 11G显存跑batch32直接OOM调到8才能跑但loss波动很大验证集精度的方差也变大。原因batch太小导致梯度估计噪声大Swin的BN层在小batch下统计量不稳定。解决显存不够时优先用混合精度训练torch.cuda.amp常见可以把batch翻倍batch仍然小就增大梯度累积步数模拟更大batch的梯度更新。另外把num_workers调到4以上pin_memory打开能减少GPU空闲等待缓解训练的不确定性。5.5 混淆矩阵里某一类故障被模型全部认成另一类现象看混淆矩阵内圈故障那行几乎全落在外圈故障列其他类正常。原因两类故障的特征频率在尺度序列覆盖的低频区间重合或者是该故障样本量过少、时频图模式太单一模型没有见过足够多样的相位和负载变化。解决检查尺度序列是否覆盖了该类故障的特征频率区间低转速下特征频率可能低于1Hzscales下限要从1改到0.5样本少就提高窗口重叠率到75%或者做数据增强给时频图加随机噪声、时间小幅平移让模型对相位变化不敏感。改完后单独看那一类的召回率有没有上升而不是整体准确率。6. 结果验证与部署边界特征可视化、特征频率对照与模型导出模型训完不代表方案结束还要确认它是靠真正的物理特征分类而不是数据集里的偶然线索。这一步也是答辩和项目汇报时最能加分的部分。6.1 t-SNE特征可视化与混淆矩阵看模型有没有学到分簇from sklearn.manifold import TSNE import matplotlib.pyplot as plt # 把模型去掉head取倒数第二层的特征向量 model.eval() features, labels [], [] with torch.no_grad(): for x, y in val_loader: x x.to(device) feat model.forward_features(x) # timm模型内置的feature方法 features.append(feat.mean(dim(1, 2)).cpu()) labels.append(y) features torch.cat(features).numpy() tsne TSNE(n_components2, perplexity30, random_state0).fit_transform(features) plt.scatter(tsne[:, 0], tsne[:, 1], ctorch.cat(labels), cmaptab10, s10) plt.savefig(tsne_check.png)forward_features输出的是Swin最后一个stage的特征图形状是(B, H/32, W/32, C)对空间维度求均值得到一个向量代表整张图的全局特征。t-SNE图里四类颜色明显分成四簇说明模型学到的是类间可分的模式如果某一类分簇散乱优先怀疑训练数据本身就不够丰富。配合混淆矩阵看每一个类的召回率比单看准确率更能暴露问题。6.2 特征频率对照表用物理规律验证诊断可信度训练结束后要做最后一个验证把模型最关注的频带和轴承故障特征频率的理论值对照。滚动轴承故障特征频率公式如下需要知道转频fr、滚动体个数n、滚动体直径d、节径D和接触角α故障位置特征频率公式典型表现时频图外圈 BPFOn/2 * fr * (1 - d/D * cosα)固定亮带无调制内圈 BPFIn/2 * fr * (1 d/D * cosα)亮带带边带有转频调制滚动体 BSFD/(2d) * fr * (1 - (d/D)^2 * cos²α)双频亮带常有间隔不均保持架 FTFfr/2 * (1 - d/D * cosα)极低频亮带对照方法是看测试集里该类样本的时频图确认图中是否有对应特征频率及其倍频的亮带。模型如果主要靠这些带分类那它的决策有明显物理依据不是纯粹在拟合数据纹理。6.3 TorchScript导出部署前的边界锁定model.eval() dummy torch.randn(1, 3, 224, 224).to(device) traced torch.jit.trace(model, dummy) traced.save(bearing_swin.pt)导出时最容易被忽略的是预处理参数必须和训练时完全一致。比如训练时用了ImageNet的mean/std归一化部署端也要用同一组数值输入尺寸固定224x224cam或者上位机送进来的图必须先缩放再归一化。我早前踩过一次训练归一化没做、部署时做了测试集准确率从96%掉到88%排查半天发现是预处理链路不一致。这套方案真正的价值不是把准确率堆到99%而是让故障诊断有从信号到图像的完整链路、有物理公式可解释、有可落地的模型导出方案在这个方向上按数据段划分、按特征频率验证、按预处理锁定这三件事养成的习惯至今都在帮我省时间。希望帮到你。本文还有配套的精品资源点击获取
