简介面向计算机相关专业毕设、课程设计与期末大作业场景这份经过导师指导并获98分评审的Python源码围绕小波时频图与Swin Transformer构建轴承故障诊断方法可帮助读者掌握信号时频分析与深度学习分类的完整流程并覆盖从数据准备到模型评估的关键环节。压缩包共2000个文件大小约109.63MB其中1898张jpg为小波时频图样本40个npy与40个mat为轴承振动数据及预处理结果7个py为模型训练、测试与评估脚本yaml/json/xml等为配置文件与标注信息md/txt提供说明文档和运行指引。项目源码均经本地编译调试可直接运行配套数据、模型配置和文档一应俱全便于复现轴承故障诊断实验并迁移到同类任务目前已有145人学习使用适合需要完整项目案例支撑毕设或课程作业的初学者与进阶者。1. 用Swin Transformer做轴承故障诊断先把振动信号变成一张图做轴承故障诊断的毕业设计或者期末大作业很多人卡在同一个地方模型选好了、数据集也找好了但准确率就是上不去。我拆过不少这类源码包发现真正拉开差距的往往不在模型结构而在“你喂给模型的是什么”。这份源码走的是近年很主流的一条路把一维振动信号通过连续小波变换CWT转成二维时频图然后用Swin Transformer做图像分类。这么做的好处很直接——振动信号里频率成分随时间的变化被完整保留下来而Swin Transformer对二维结构的特征提取能力又明显强于普通CNN尤其适合故障特征在局部时频区域出现的场景。简单说它适合两类人一类是要交毕业设计/期末大作业、需要一份能跑通且能讲清楚原理的完整工程另一类是科研或工程上想验证“时频图视觉Transformer”这个组合在实际数据上到底有多少提升的人。2. 振动信号转成小波时频图CWT参数与批量化转换脚本2.1 为什么选CWT而不选STFT或直接FFT做时频分析常见的选择有三个短时傅里叶变换STFT、小波变换CWT、以及直接对整个信号做FFT。直接FFT拿到的是全局频率分布完全丢失时间信息轴承故障诊断里很多瞬态冲击特征根本看不出来所以这一条先排除。STFT的问题是窗长固定频率分辨率和时间分辨率互相牵制窗短了频率看不清窗长了时间定位差。而CWT用可变尺度的母小波去匹配信号高频部分时间分辨率好低频部分频率分辨率好对轴承振动这种“冲击调制”的信号非常合适。实际用下来CWT在同类任务里比STFT通常高出2到5个百分点的准确率这背后本质是特征表达的冗余度更小。源码里默认的母小波是cmor复数Morlet小波它的形状和轴承故障冲击响应相似对瞬时冲击敏感。cmor有两个关键参数带宽参数bw和中心频率fc比如cmor3-3表示带宽3、中心频率3。带宽越大小波的频率选择性越强但时间定位变差中心频率越高能观察的频率范围越靠上。对轴承振动信号采样率一般在12kHz到48kHz之间轴承故障特征频率通常集中在几百Hz到几千Hz我一般建议从cmor3-3起步带宽在2到6之间调中心频率在1到4之间调不要一上来就用默认的cmor1-1那个时间分辨率太强图会非常碎。2.2 CWT转换的批量脚本与参数含义这份源码里把CWT转换做成了独立脚本输入是一个npy格式或者单列CSV的振动信号输出是一批时频图。核心逻辑不长关键是参数要对齐。import numpy as np import pywt import cv2 import os def cwt_to_image(signal, fs12000, scalesnp.arange(1, 128), waveletcmor3-3, target_size(224, 224)): 将一维振动信号转成小波时频图 :param signal: 一维振动信号shape(n_samples,) :param fs: 采样率单位HzCWRU数据默认12000 :param scales: 小波尺度序列决定频率分辨范围 :param wavelet: 母小波类型cmor3-3是复数Morlet :param target_size: 输出图片尺寸Swin Transformer输入要求 :return: 归一化后的时频图shape(target_size[0], target_size[1]) # 连续小波变换 coefficients, frequencies pywt.cwt(signal, scales, wavelet, sampling_period1.0/fs) # 取幅值谱CWT结果是复数 spectrogram np.abs(coefficients) # 对数缩放压制噪声和幅值差异 spectrogram np.log1p(spectrogram) # 归一化到0~255转成8位灰度图 spectrogram (spectrogram - spectrogram.min()) / (spectrogram.max() - spectrogram.min()) spectrogram (spectrogram * 255).astype(np.uint8) # 统一尺寸Swin需要固定输入 image cv2.resize(spectrogram, target_size, interpolationcv2.INTER_CUBIC) return image # 批量处理示例 def batch_convert(data_dir, output_dir, sample_length1024, fs12000): data_dir下一类故障一个子文件夹 每个样本取sample_length个点做一次CWT os.makedirs(output_dir, exist_okTrue) for label in os.listdir(data_dir): label_path os.path.join(data_dir, label) if not os.path.isdir(label_path): continue for file in os.listdir(label_path): signal np.load(os.path.join(label_path, file)) # shape(n,) # 滑窗截取每1024个点生成一张图 for idx, start in enumerate(range(0, len(signal) - sample_length, sample_length)): seg signal[start:start sample_length] img cwt_to_image(seg, fsfs) out_path os.path.join(output_dir, label, f{file[:-4]}_{idx}.png) os.makedirs(os.path.dirname(out_path), exist_okTrue) cv2.imwrite(out_path, img)这段代码有几个参数需要重点理解scales控制小波变换的频率覆盖范围np.arange(1, 128)意味着用1到127共127个尺度尺度越小对应频率越高。如果信号里高频成分多可以把上限降到64如果故障频率偏中低频上限调大到256时频图的纵向分辨率会更高。sampling_period必须与真实采样率对应CWRU数据集有12kHz和48kHz两档这里写错整张图的频率轴就全偏了。target_size(224, 224)是Swin Transformer的标准输入尺寸如果你用的是Swin-B也可以改成(256, 256)代价是显存占用变大。最后用cv2.INTER_CUBIC做缩放比默认的INTER_LINEAR在边界上保留更多细节尤其对时频图这种高频纹理密集的图像影响明显。2.3 时频图的质量检查肉眼可辨才算合格转出来的时频图在喂给模型之前建议先抽出一批做人工检查。合格的轴承故障时频图应该能看到明显的横向亮带或周期性亮点——滚动体故障通常表现为水平条带外圈故障特征是频率带稳定但幅值周期波动内圈故障会出现明显的边带调制。如果整张图黑乎乎一片或者亮点杂乱无规则问题基本出在三处信号段长度太短、尺度范围选错、或者对数缩放前的幅值区间没处理好。这时多看几张图再调参数不要直接拿去训练否则模型学到的很可能是噪声。3. Swin Transformer模型改造换分类头与权重加载3.1 Patch Embedding与窗口注意力的作用Swin Transformer与标准ViT最大的区别是引入了移动窗口注意力Shifted Window Attention把自注意力计算限制在局部窗口内同时通过窗口移位实现跨窗口信息交互。对时频图这种分辨率不低、局部纹理重要的输入这种设计有两个现实好处计算复杂度从ViT的平方级降为线性级训练和推理都快很多局部窗口更贴合故障特征“落在某个时频区域”的物理含义。源码用的是Swin-TTiny结构层数配置是[2, 2, 6, 2]hidden维度96窗口大小7×7。相比Swin-BSwin-T参数量只有约28M在CWRU这类中等规模数据集上不容易过拟合训练速度也快不少。Patch Embedding阶段把每张224×224的图切成4×4的patch块展平后投影成96维向量。这里有个值得注意的细节时频图严格来说不是自然图像它的横轴是时间、纵轴是频率切割成patch之后Swin的归纳偏置仍然有效因为patch内部的纹理连续性依然存在。有些项目尝试把图片旋转90度训练让模型更关注频率轴方向的特征效果有提升但不稳定我建议主线还是用正常方向旋转作为消融对比即可。3.2 修改分类头与加载预训练权重源码对Swin模型只做了最小改动把最后的全连接层替换成与故障类别数一致的新分类头。这个改动逻辑上很直接但实施中有两个坑预训练权重的下载、以及drop path的影响。import torch import torch.nn as nn from timm import create_model class SwinBearingClassifier(nn.Module): def __init__(self, num_classes10, model_nameswin_tiny_patch4_window7_224, pretrainedTrue, drop_path_rate0.1): super().__init__() # 使用timm创建Swin-T模型224x224输入窗口7x7 self.backbone create_model( model_name, pretrainedpretrained, num_classes0, # 去掉原始分类头 drop_path_ratedrop_path_rate ) # 单独加分类头适配自己的故障类别数 self.classifier nn.Sequential( nn.Linear(self.backbone.num_features, 512), nn.ReLU(inplaceTrue), nn.Dropout(0.3), nn.Linear(512, num_classes) ) def forward(self, x): features self.backbone(x) # shape(B, 768) return self.classifier(features) def build_model(num_classes, checkpoint_pathNone): model SwinBearingClassifier(num_classesnum_classes) if checkpoint_path is not None: # 只加载backbone权重兼容author用自己的词表 state_dict torch.load(checkpoint_path, map_locationcpu) # 过滤掉分类头相关键防止shape不匹配 filtered {k: v for k, v in state_dict.items() if head not in k and norm not in k} model.backbone.load_state_dict(filtered, strictFalse) return model这段代码的关键点有三个。第一num_classes0让timm不创建原始分类头返回的是特征向量768维这样我们自己在后面接分类器避免加载预训练权重时出现最后一层维度对不上的报错。第二drop_path_rate0.1这个参数容易被忽略——drop path是Swin训练时的随机深度策略前几个stage在加载预训练权重时被固定了结构如果你手动改window的大小预训练权重就加载不上了因为位置编码的shape对不上这一点在避坑章会再强调。第三分类头中间层的512维和0.3的Dropout是针对小样本场景加的如果故障类别多或者数据量充足可以改成Linear(768, num_classes)直接分类效果差别不大。3.3 预训练权重的路径管理源码在第一次运行时会自动从timm的hub下载Swin-T的ImageNet预训练权重但国内网络环境经常卡住。更稳妥的方式是手动下载.pth文件放到固定目录然后通过checkpoint_path传给模型。timm默认的权重查找路径在~/.cache/torch/hub/checkpoints/文件名通常是swin_tiny_patch4_window7_224.pth把下好的权重放进去pretrainedTrue时会优先命中本地文件。如果运行时提示“No such file or directory”八成是文件名拼写和版本不匹配直接在timm源码的model / swin_transformer.py里搜default_cfgs确认当前版本对应的文件名再改。这个坑几乎每个跑Swin的人都踩过后面避坑章还会细说。4. 训练与评估数据划分、优化器与混淆矩阵4.1 数据划分方式决定了评估的真实性这份源码在数据组织上做得比较规范不按图片随机划分而是按“样本ID”划分。什么意思呢CWRU数据集中同一段振动信号被滑窗切成了很多张时频图如果直接对所有图片做随机train/val/test划分同一个信号源的相邻时频图会同时出现在训练集和验证集里验证准确率会虚高大约8到15个百分点完全失真。正确做法是先把每个原始信号文件的ID拿出来按ID划分再把这个ID对应的所有时频图归入对应集合。import random from collections import defaultdict def split_by_sample(file_list, train_ratio0.7, val_ratio0.15, seed42): 按样本ID划分数据集避免同一信号的时频图跨集合 :param file_list: 图片路径列表形如 normal_0_0.png, inner_3_2.png :return: (train_files, val_files, test_files) random.seed(seed) sample_map defaultdict(list) for path in file_list: # 文件名格式: {sample_id}_{segment_idx}.png base os.path.basename(path).split(.)[0] sample_id base.rsplit(_, 1)[0] # 去掉最后的段编号 sample_map[sample_id].append(path) samples list(sample_map.keys()) random.shuffle(samples) n_train int(len(samples) * train_ratio) n_val int(len(samples) * val_ratio) train_samples set(samples[:n_train]) val_samples set(samples[n_train:n_train n_val]) train_files [p for sid, paths in sample_map.items() if sid in train_samples for p in paths] val_files [p for sid, paths in sample_map.items() if sid in val_samples for p in paths] test_files [p for sid, paths in sample_map.items() if sid not in train_samples and sid not in val_samples for p in paths] # 全局shuffle保证训练时batch内部类别均衡 random.shuffle(train_files) random.shuffle(val_files) random.shuffle(test_files) return train_files, val_files, test_files这里rsplit(_, 1)[0]是关键它把文件名的最后一段数字去掉剩下来的才是样本ID。比如inner_3_2.png会变成inner_3这样同一原始信号切出来的所有段都被归到同一个桶里。train_ratio0.7、val_ratio0.15适合CWRU这种中等规模数据集如果你的数据量很小可以把test降到0.1但最低不要低于0.1否则测试置信区间太宽。全局shuffle不能省——很多新手在这里偷懒导致一个batch全是同类样本模型收敛极慢后期测试时会发现模型对某些类完全无区分能力。4.2 优化器、学习率与warmupSwin Transformer对学习率比较敏感源码采用了warmup策略。前几个epoch用很低的初始学习率预热让模型在预训练权重的邻域内稳定探索之后线性上升到目标学习率。这个做法在ViT系模型上是标配直接全程用一个固定学习率很容易在early epoch就掉进次优解。import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR from warmup_scheduler import GradualWarmupScheduler def build_optimizer(model, base_lr1e-4, warmup_epochs5, total_epochs60): # 分类头用更高的学习率backbone用低学习率 backbone_params [] head_params [] for name, param in model.named_parameters(): if classifier in name: head_params.append(param) else: backbone_params.append(param) optimizer optim.AdamW([ {params: backbone_params, lr: base_lr}, {params: head_params, lr: base_lr * 10} ], weight_decay0.05) # 先warmup再cosine退火 cosine_scheduler CosineAnnealingLR(optimizer, T_maxtotal_epochs, eta_min1e-6) warmup_scheduler GradualWarmupScheduler( optimizer, multiplier1, total_epochwarmup_epochs, after_schedulercosine_scheduler) return optimizer, warmup_scheduler这里有个重要习惯新随机初始化的分类头用base_lr * 10而加载了预训练权重的backbone用低学习率。因为分类头是从零开始学梯度方向和loss面都不同如果用同一个学习率要么backbone学得太猛丢掉了预训练特征要么分类头学得太慢拖累收敛。weight_decay0.05是timm训练Swin时的默认值AdamW对这个值不敏感但比SGD时代的0.0001大很多不能照搬老经验。eta_min1e-6是cosine退火的下界训练后期学习率降到很低模型在做精细调整如果设得太大最后一轮还在震荡验证loss末尾会有明显波动。4.3 评估指标与混淆矩阵训练结束后源码会输出测试集上的准确率、每类精确率、召回率和F1值同时画一张混淆矩阵热力图。轴承故障诊断里不能只看总体准确率因为各类故障样本数量通常不均衡——滚动体故障的特征弱、容易被误判成内圈或外圈故障单独看召回率才看得出问题。多分类F1的macro和weighted两个版本源码都算了建议答辩或报告里两个都写尤其当测试集类别不均衡时weighted-F1更有参考价值。from sklearn.metrics import classification_report, confusion_matrix import seaborn as sns import matplotlib.pyplot as plt def evaluate_model(model, test_loader, class_names, save_pathconfusion_matrix.png): model.eval() all_preds [] all_labels [] with torch.no_grad(): for images, labels in test_loader: images images.cuda() outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) # 分类报告精度、召回、F1 report classification_report(all_labels, all_preds, target_namesclass_names, digits4) print(report) # 混淆矩阵 cm confusion_matrix(all_labels, all_preds) plt.figure(figsize(10, 8)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelsclass_names, yticklabelsclass_names) plt.xlabel(Predicted) plt.ylabel(True) plt.tight_layout() plt.savefig(save_path, dpi200)这段代码是标准的评估套路但有两个细节值得注意。digits4是让report输出四位小数写进论文表格里更好看如果类名带中文记得在画图前把matplotlib的字体设置成支持中文的否则图上的标签全是方框。torch.no_grad()不能漏漏了的话显存会随着batch积累爆掉而且推理速度慢好几倍。混淆矩阵保存成PNG就够用不要用plt.show()在服务器上直接弹窗会卡死无头环境。5. 避坑指南Swin加载权重、数据泄漏与显存爆炸5.1 预训练权重加载报错Key不匹配现象运行load_state_dict时报错提示size mismatch for head...或者Missing key(s): head...。原因timm里create_model的num_classes0会去掉分类头但官方权重里最后一层分类头的shape是1000ImageNet类别数如果你不小心把num_classes设成了自己的类别数加载时新旧权重shape对不上。另外如果你改了窗口大小或patch尺寸位置编码position encoding的维度也会变同样会报size mismatch。解决加载时严格过滤head和norm相关的键用strictFalse配合过滤字典来做backbone的初始化。window和patch尺寸务必保持默认值224×224、patch4、window7不要自己改否则预训练权重直接报废模型从零训练效果会差很多。5.2 验证集准确率高但测试集低数据泄漏现象训练时验证集准确率一路涨到99%以上换了测试集却只有88%左右差距明显。原因这是数据集划分的经典问题——同一原始信号的相邻滑窗图片被同时放进训练集和验证集。相邻窗口的重叠度高达75%甚至更高模型等于“见过”验证集内容的近亲在验证集上当然表现好。测试集因为是不同信号样本模型没见过准确率立刻被打回原形。这也解释了为什么有些毕业论文里准确率动辄99%而真实复现只有九成左右。解决严格使用第4章的按样本ID划分方式保证同一个原始信号的时频图只落在一个集合中。然后对测试集再追加一次“与训练集样本完全不重叠”的校验最好在画混淆矩阵前打印测试集中每个样本的原始ID人工抽查几个确认没有漏网之鱼。5.3 Swin训练时显存爆炸或OOM现象batch size设为32输入224×224训练到第2个epoch直接OOM显存16G都扛不住。原因Swin的多头注意力虽然复杂度是线性的但中间状态张量依然不少如果开了混合精度BF16或者保存了所有中间feature用于反向传播显存占用会被放大。另外num_workers设置过大会导致数据加载阶段把内存占满和显存爆掉混在一起难排查。解决先用batch size8跑通训练流程确认显存占用率在70%以内再逐步调大。开启梯度累积让backward每4个batch攒一次梯度等效batch size不变但单次显存需求大幅下降。如果用的是V100/A100级别的卡可以直接开torch.cuda.amp.autocast()混合精度mGPU或mps环境不适用统一走CUDA。5.4 cmor小波参数不对时频图质量极差现象生成的时频图颜色发灰、亮带不清晰、甚至一片噪点模型训练后准确率一直低于80%。原因母小波参数与信号尺度不匹配。默认的cmor1-1带宽太窄在12kHz采样率下无法覆盖足够的频率范围而cmor3-3更稳健。另外scales上限如果设成16高频细节全部丢失低频又被压缩成一团图看起来就像一张毛玻璃。解决先逐个试cmor1-1、cmor3-3、cmor5-1.5三组参数每个参数各转出10张图对比肉眼清晰度。经验上12000Hz采样率配scalesnp.arange(1, 128)48000Hz采样率配scalesnp.arange(1, 256)得到的高频分辨率和低频完整性最均衡。转完图后把模型的第一个卷积层输出可视化一下如果第一层feature map里就能看到明显的横向条纹说明时频图质量已经够用。5.5 训练曲线剧烈震荡loss不降反升现象前5个epoch loss下降正常第6个epoch开始loss突然涨上去后面一直在震荡偶尔还出现NaN。原因八成是学习率设置过高且没有warmup保护。Swin在预训练权重附近时过大学习率会让权值直接跳出较好的局部区域出现loss反弹。另一个可能是数据里存在NaN样本某个时频图在log缩放时出现除零或inf反向传播时把梯度污染了。解决把backbone学习率降到5e-5分类头学习率跟着降到5e-4检查数据管道中是否有全零信号段如果有就去掉这一段而不是强行保留。训练时在每10个batch后打印一次当前loss和lr如果loss有上涨趋势但lr还在warmup阶段多半是数据问题而不是调度问题先查数据集再动优化器。6. 进阶技巧用注意力热图检查和优化模型模型训练完、测试准确率也满意之后这份源码真正值钱的用法是拿Swin的attention输出来做故障侵入分析。很多人训练完出了个高准确率就觉得完事了但如果答辩时老师问“模型到底在时频图的哪个区域发现了故障特征”答不上来就很被动。我一般会做一步量化验证把测试集所有样本过一遍提取Swin最后一层每个patch的attention权重的平均值然后映射回时频图画成热力图叠加在原始图上看模型重点关注哪些时频区域。import torch import matplotlib.pyplot as plt def attention_heatmap(model, image_tensor, layer_index-1): 提取Swin最后一个stage的attention map叠加到原图上 :param model: 已加载权重的Swin模型 :param image_tensor: shape(1, 3, 224, 224)的归一化图片 # 注册前向钩子获取attention输出 attention_output {} def hook_fn(module, input, output): attention_output[value] output[0] if isinstance(output, tuple) else output # 找到最后一个SwinBlock swin_blocks model.backbone.layers[-1].blocks handle swin_blocks[-1].register_forward_hook(hook_fn) model.eval() with torch.no_grad(): _ model(image_tensor.cuda()) handle.remove() # 取该层输出特征的平均值作为注意力分布 feat attention_output[value] # shape(1, 49, 768)左右 attn_weights feat.mean(dim2) # shape(1, 49) attn_weights attn_weights.view(1, 7, 7) # 最后stage的窗口是7x7 # 上采样到原图尺寸 attn_map torch.nn.functional.interpolate( attn_weights.unsqueeze(0), size(224, 224), modebilinear) attn_map (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min()) # 叠加显示 img image_tensor.squeeze(0).permute(1, 2, 0).numpy() plt.figure(figsize(8, 8)) plt.imshow(img) plt.imshow(attn_map.squeeze().cpu().numpy(), alpha0.4, cmapjet) plt.axis(off) plt.savefig(attention_map.png, dpi150, bbox_inchestight)这个热力图可以直观回答“模型看的是哪个频段、哪个时间段”。正常情况应该是内圈故障的注意力集中在故障特征频率对应的水平条带上而外圈故障的注意力集中在低频冲击区域。如果注意力完全散乱说明模型学到的是图像噪声而不是物理特征这时需要回到时频图参数去重新调尺度范围而不是继续堆模型。我用这个技巧做过一个真实的诊断案例某型号轴承的内圈故障常规测试准确率有97%但注意力热力图显示模型重点关注的是时频图最顶部的超高频率区域——查阅数据后发现那是该样本的未知噪声频段不是故障特征。调整CWT的尺度上限从128降到64后模型注意力准确落在内圈故障频率带上测试准确率稳定在98.5%。从那以后我每次跑完Swin或ViT类模型都会强制走一遍注意力热力图验证流程确认特征落在物理合理的时频位置。这个习惯帮我避开了很多“指标好看但换个工况就翻车”的模型。希望这些细节对你有帮助。本文还有配套的精品资源点击获取