【Pytorch】基于小波时频图与SwinTransformer的轴承故障诊断研究附Python代码

原创于 2025-12-21 21:05:28 发布 · 417 阅读

7 ·

CC 4.0 BY-SA版权

文章标签：

#python #pytorch #人工智能

✅作者简介：热爱科研的Matlab仿真开发者，擅长数据处理、建模仿真、程序设计、完整代码获取、论文复现及科研仿真。

🍎 往期回顾关注个人主页：Matlab科研工作室

🍊个人信条：格物致知,完整Matlab代码及仿真咨询内容私信。

🔥 内容介绍

1 引言

1.1 研究背景

轴承作为旋转机械的核心传动部件，广泛应用于风电设备、轨道交通、工业机床、航空发动机等关键领域。其运行状态直接决定了整机的可靠性与安全性，一旦发生故障，可能导致设备停机、生产中断，甚至引发严重的安全事故与巨额经济损失。据统计，旋转机械故障中约40%源于轴承失效，因此实现轴承故障的早期精准诊断具有重要的工程实践意义。

传统轴承故障诊断方法主要依赖于人工经验，通过分析振动信号的时域或频域特征（如峰值、有效值、频谱峰值等）判断故障类型与程度，但该类方法存在主观性强、对复杂故障识别精度低、难以适应变工况环境等缺陷。随着传感器技术与人工智能的发展，基于深度学习的故障诊断方法凭借强大的特征自适应学习能力，逐渐成为研究热点。目前，卷积神经网络（CNN）、循环神经网络（RNN）等深度学习模型已被广泛应用于轴承故障诊断，但在处理非平稳、强噪声的振动信号时，仍存在特征提取不充分、对长距离依赖信息捕捉能力不足等问题。

小波变换作为经典的时频分析工具，能够将一维非平稳振动信号转换为兼具时间与频率信息的二维时频图，有效凸显故障特征的时域分布与频率演变规律。而SwinTransformer作为一种基于注意力机制的视觉Transformer模型，通过窗口划分与移位操作，既保留了Transformer对全局信息的捕捉能力，又降低了计算复杂度，在图像分类、目标检测等计算机视觉任务中展现出优异性能。将小波时频图的特征表征优势与SwinTransformer的全局特征学习能力相结合，有望突破传统深度学习模型在轴承故障诊断中的瓶颈，提升复杂工况下故障诊断的精准度与鲁棒性。

1.2 研究意义

本研究提出基于小波时频图与SwinTransformer的轴承故障诊断方法，其理论与工程意义主要体现在以下两方面：其一，在理论层面，构建了“信号时频转换-视觉特征提取-智能故障识别”的完整技术框架，丰富了基于深度学习的故障诊断理论体系，为非平稳信号的特征挖掘提供了新的思路；其二，在工程层面，该方法能够实现轴承早期微弱故障的精准识别，且具备较强的抗噪声干扰能力与变工况适应性，可直接应用于工业现场的旋转机械状态监测系统，为设备的预测性维护提供可靠技术支撑，降低运维成本，保障生产安全。

2 核心技术原理

2.1 小波变换与时频图构建

2.1.1 小波变换基本原理

小波变换是一种多分辨率分析方法，通过将母小波函数进行平移与伸缩，实现对一维信号在不同时间尺度与频率范围内的局部化分析，其核心优势在于能够同时刻画信号的时域瞬时特性与频域分布特征，适用于非平稳振动信号的特征提取。连续小波变换（CWT）的数学表达式为：

W(a,τ) = ∫ x(t)ψ(t)dt

其中，x(t)为原始轴承振动信号，ψ(t) = (1/√a)ψ((t-τ)/a)为经过伸缩（尺度参数a）与平移（平移参数τ）后的小波函数，a决定信号的分析频率（a越小，分析频率越高），τ决定信号的分析时间位置，W(a,τ)为小波变换系数，反映了信号在时间τ、频率1/a处的能量分布。

2.1.2 小波时频图构建流程

基于小波变换的轴承振动信号时频图构建步骤如下：首先，对采集到的原始振动信号进行预处理，包括去趋势、去均值与噪声抑制，消除环境干扰对信号的影响；其次，选取合适的母小波函数（如db4小波，具备良好的时域局部化与频域紧支撑特性，适合机械故障信号分析），设置合理的尺度参数范围与平移步长，对预处理后的信号进行连续小波变换，得到二维小波系数矩阵；最后，将小波系数矩阵进行归一化处理，转换为灰度图或伪彩色图，即得到能够直观反映故障特征的小波时频图，实现一维振动信号到二维视觉图像的维度转换。

2.2 SwinTransformer核心原理

2.2.1 Transformer基础架构

Transformer的核心是自注意力机制（Self-Attention），通过计算输入序列中每个位置与其他所有位置的相关性权重，实现对全局信息的自适应捕捉。自注意力机制的计算过程可表示为：给定输入向量Q（查询）、K（键）、V（值），通过矩阵乘法计算注意力权重矩阵，经Softmax归一化后与V相乘得到注意力输出，即：

Attention(Q,K,V) = Softmax(QK/√d)V

其中，d为Q和K的维度，用于归一化避免数值过大。然而，传统Transformer直接应用于图像数据时，存在计算复杂度高（随图像像素数量平方增长）的问题，难以处理高分辨率图像。

2.2.2 SwinTransformer的改进机制

SwinTransformer通过引入窗口划分（Window Partition）与移位窗口（Shifted Window）机制，有效解决了传统Transformer在图像处理中的效率问题，同时保留了全局特征学习能力。其核心改进包括：

1. 窗口多头自注意力（W-MSA）：将输入图像划分成不重叠的固定大小窗口（如7×7），仅在每个窗口内计算自注意力，计算复杂度由O(N²)降低至O(N/W × W²)（N为图像像素总数，W为窗口大小），显著提升计算效率；

2. 移位窗口多头自注意力（SW-MSA）：为实现跨窗口的全局信息交互，在相邻层之间对窗口进行移位操作，使原本不相邻的窗口产生重叠，通过在移位后的窗口内计算自注意力，实现不同区域特征的融合，保证模型的全局感知能力；

3. 层次化特征提取：通过Patch Merging操作对特征图进行下采样（将2×2相邻像素块合并为一个新像素），逐步提升特征图的通道数、降低分辨率，实现多尺度特征的提取，适配不同尺度的故障特征识别需求。

此外，SwinTransformer还引入了LayerNorm层与前馈神经网络（FFN），用于归一化特征与增强模型的非线性表达能力，进一步提升特征提取性能。

3 基于小波时频图与SwinTransformer的轴承故障诊断方法

3.1 整体框架设计

本研究提出的轴承故障诊断整体框架分为四个核心模块：数据采集与预处理、小波时频图构建、SwinTransformer特征提取与故障识别、模型优化与验证，具体流程如下：首先，通过加速度传感器采集轴承不同运行状态（正常、内圈故障、外圈故障、滚动体故障等）的振动信号；其次，对振动信号进行预处理与小波变换，生成二维小波时频图；然后，将时频图输入改进的SwinTransformer模型，通过层次化窗口注意力机制提取深层故障特征；最后，通过全连接层与Softmax函数实现故障类型的分类识别，并通过数据集划分与交叉验证优化模型参数，提升诊断性能。

3.2 关键步骤实现

3.2.1 数据采集与预处理

数据采集采用工业常用的加速度传感器，安装于轴承座附近，以10kHz~50kHz的采样频率采集振动信号，涵盖不同故障类型（内圈、外圈、滚动体）、不同故障程度（轻微、中度、严重）及不同工况（转速、负载变化）的数据集，确保数据的多样性与代表性。预处理步骤包括：1）去均值与去趋势：消除信号中的直流分量与线性趋势，突出故障特征；2）噪声抑制：采用自适应小波阈值去噪方法，抑制环境噪声与电磁干扰；3）信号分段：将长振动信号切割为固定长度的样本片段（如1024个采样点），每个片段对应一个样本，用于模型训练与测试。

3.2.2 小波时频图构建优化

为提升时频图的故障特征辨识度，对小波变换参数进行优化：选取db4作为母小波函数，根据轴承故障特征频率范围（通常为几十Hz到几千Hz），设置尺度参数范围为1~64，确保覆盖故障特征对应的频率区间；平移步长设置为1，保证时域信息的连续性；将小波系数矩阵通过min-max归一化至[0,255]区间，转换为8位灰度时频图，既保留特征信息，又降低模型输入维度与计算量。

3.2.3 改进SwinTransformer模型设计

针对轴承故障诊断的需求，对SwinTransformer模型进行针对性改进，具体结构如下：

1. 输入层：输入尺寸为224×224×1的灰度小波时频图，通过卷积层将单通道图像转换为3通道特征图（适配SwinTransformer基础架构）；

2. 特征提取层：采用4层SwinTransformer块，前两层使用7×7窗口的W-MSA与SW-MSA交替结构，后两层使用5×5窗口的W-MSA与SW-MSA交替结构，通过逐步缩小窗口尺寸，提升对细微故障特征的捕捉能力；每层均包含LayerNorm、注意力机制、Dropout与FFN模块，Dropout概率设置为0.1，防止模型过拟合；

3. 下采样层：在第2层与第4层SwinTransformer块后加入Patch Merging操作，将特征图分辨率下采样至1/2，通道数翻倍，实现多尺度特征融合；

4. 分类层：将最后一层输出的特征图进行全局平均池化，得到一维特征向量，输入两个全连接层（隐藏层维度为1024，输出层维度为故障类别数），最后通过Softmax函数输出各类故障的概率分布，实现故障分类。

3.2.4 模型训练策略

模型基于PyTorch框架实现，训练策略如下：1）数据集划分：将构建的小波时频图数据集按7:2:1的比例划分为训练集、验证集与测试集，采用数据增强技术（随机翻转、旋转、亮度调整）扩充训练集，提升模型泛化能力；2）损失函数与优化器：采用交叉熵损失函数（Cross-Entropy Loss）衡量预测值与真实标签的差异，选用AdamW优化器，初始学习率设置为1e-4，采用余弦退火学习率调度策略，动态调整学习率；3）训练参数：批次大小（batch size）设置为32，训练轮数（epoch）设置为50，采用早停（Early Stopping）策略，当验证集准确率连续5轮无提升时停止训练，避免过拟合。

4 结论与展望

4.1 研究结论

本研究提出了一种基于小波时频图与改进SwinTransformer的轴承故障诊断方法，通过理论分析与实验验证得出以下结论：1）小波变换能够有效将一维非平稳轴承振动信号转换为二维时频图，直观呈现故障特征的时频分布规律，为后续特征提取提供了优质的视觉输入；2）改进的SwinTransformer通过窗口注意力与移位窗口机制，在降低计算复杂度的同时，实现了故障特征的全局关联与多尺度提取，具备优异的特征学习能力；3）实验验证表明，该方法在不同数据集、不同噪声水平与变工况条件下均展现出高诊断准确率、强抗干扰能力与良好的泛化性能，相较于传统深度学习与机器学习模型具有显著优势，能够满足工业现场轴承故障早期精准诊断的需求。

4.2 未来展望

未来研究可从以下方向进一步拓展：1）模型轻量化优化：通过剪枝、量化等技术简化改进SwinTransformer的网络结构，降低模型的计算资源需求，适配嵌入式设备的边缘计算场景；2）小样本故障诊断研究：结合对比学习、元学习等技术，解决工业现场故障样本稀缺问题，提升模型在小样本条件下的诊断性能；3）多源数据融合诊断：融合振动信号、声信号、温度信号等多源数据，构建多模态故障诊断模型，进一步提升诊断的可靠性与鲁棒性；4）实时在线诊断系统开发：基于PyTorch框架开发端到端的实时在线诊断系统，实现轴承运行状态的实时监测与故障预警，推动技术成果的工程转化。