基于Swin Transformer的轴承故障诊断方法

1. 项目背景与核心思路

轴承作为旋转机械的核心部件，其故障诊断一直是工业设备健康管理的关键课题。传统振动信号分析方法存在特征提取依赖专家经验、诊断准确率受噪声影响大等痛点。我们团队尝试将时频分析与视觉Transformer结合，提出了一种基于连续小波变换(CWT)时频图和Swin Transformer的端到端故障诊断方案。

这个方案的创新点在于：

• 利用CWT将一维振动信号转换为二维时频图像，完整保留时频域特征
• 采用Swin Transformer作为特征提取主干网络，其窗口注意力机制特别适合处理具有局部相关性的时频图像
• 整个流程实现了从原始振动信号到故障类别的端到端识别，无需人工设计特征

2. 技术实现详解

2.1 数据准备与预处理

我们使用凯斯西储大学(CWRU)轴承数据集进行实验，该数据集包含正常状态和多种故障类型(内圈、外圈、滚动体故障)在不同损伤程度下的振动信号。数据预处理流程如下：

1. 信号分段：将长时序信号切分为固定长度片段(如1024个采样点)
2. 标准化处理：对每个信号片段进行z-score标准化
3. 数据增强：通过添加高斯噪声、时间偏移等方式扩充训练集

python复制import numpy as np
from scipy import signal

def preprocess_data(raw_signal, segment_length=1024):
    # 信号分段
    segments = []
    for i in range(0, len(raw_signal)-segment_length, segment_length//2):
        segment = raw_signal[i:i+segment_length]
        # 标准化
        segment = (segment - np.mean(segment)) / np.std(segment)
        segments.append(segment)
    return np.array(segments)

2.2 小波时频图生成

连续小波变换能将信号从时域转换到时频域，我们选用Morlet小波作为母小波，因其在时频分辨率上具有良好的平衡。关键参数设置：

• 尺度参数：根据信号采样频率和感兴趣频带确定
• 小波中心频率：典型值取6.0
• 时频图尺寸：统一调整为224×224以适应Swin Transformer输入

python复制import pywt

def generate_cwt_images(signal_segments, scales=np.arange(1, 128)):
    cwt_images = []
    for seg in signal_segments:
        coefficients, _ = pywt.cwt(seg, scales, 'morl')
        # 转换为RGB图像并调整尺寸
        img = normalize_and_resize(coefficients)
        cwt_images.append(img)
    return np.array(cwt_images)

注意：小波尺度的选择直接影响时频图质量，建议先通过功率谱分析确定信号主要频带范围，再据此调整尺度参数。

2.3 Swin Transformer模型构建

我们基于Pytorch实现了一个轻量级Swin-T变体，主要结构如下：

1. Patch Partition：将224×224图像分割为4×4的patch
2. 阶段式Transformer块：
   • 阶段1：4个Swin Transformer块，窗口大小7×7
   • 阶段2：8个块，窗口大小7×7
   • 阶段3：6个块，窗口大小14×14
3. 分类头：全局平均池化 + 全连接层

python复制import torch
from torch import nn
from swin_transformer_pytorch import SwinTransformer

model = SwinTransformer(
    hidden_dim=96,
    layers=(2, 2, 6, 2),
    heads=(3, 6, 12, 24),
    channels=3,
    num_classes=10,
    head_dim=32,
    window_size=7,
    downscaling_factors=(4, 2, 2, 2)
)

2.4 模型训练策略

采用分阶段训练策略提升模型性能：

1. 预训练阶段：

   • 使用ImageNet预训练权重初始化
   • 冻结除分类头外的所有层
   • 学习率1e-4，训练20个epoch

2. 微调阶段：

   • 解冻所有层
   • 初始学习率5e-5，余弦退火调度
   • 标签平滑(0.1)和MixUp数据增强
   • 早停策略(patience=10)

python复制from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

optimizer = AdamW(model.parameters(), lr=5e-5, weight_decay=0.05)
scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

3. 实验结果与分析

在CWRU数据集上的测试结果表明：

关键发现：

1. 时频图像比原始信号更有利于特征学习
2. Swin Transformer的窗口注意力机制能有效捕捉故障特征
3. 模型在强噪声环境下(SNR<10dB)仍保持90%+准确率

4. 工程实践建议

1. 部署优化：

   • 使用TorchScript导出模型提升推理速度
   • 量化模型减小存储占用
   • 开发基于Gradio的快速演示界面

2. 故障诊断系统设计：

   • 实时信号采集模块
   • 滑动窗口处理流水线
   • 结果可视化与报警机制

python复制# 模型量化示例
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8
)
torch.jit.save(torch.jit.script(quantized_model), "fault_diagnosis.pt")

5. 常见问题排查

1. 时频图出现伪影：

   • 检查小波尺度范围是否合适
   • 验证信号预处理是否去除直流分量
   • 尝试不同的小波基函数

2. 模型训练不收敛：

   • 检查数据标准化是否正确
   • 尝试减小学习率
   • 验证标签是否正确编码

3. 推理结果不稳定：

   • 增加滑动窗口重叠率
   • 添加结果平滑滤波
   • 检查信号采集是否受到干扰

实际部署中，我们发现轴承安装位置的不同会导致信号特征变化，建议针对不同设备分别收集数据微调模型。对于关键设备，可以采用在线学习策略持续优化模型性能。