基于STFT时频图与视觉顶会网络SwinTransformer的轴承故障诊断

1.数据预处理

傅立叶变换（Fourier Transform）用于将一个信号从时间域转换到频率域，然而它假设信号是平稳的，这意味着信号的频率成分在整个时间上是不变的。对于非平稳信号，傅立叶变换无法提供信号随时间变化的频率信息。

短时傅立叶变换通过将信号分割成多个短的、重叠的时间窗口，每个窗口内的信号近似为平稳的，然后对每个窗口分别进行傅立叶变换，从而获得信号在不同时刻的频率信息。

将CWRU数据集进行短时傅里叶变换

滚动体故障

内圈故障

正常信号

外圈故障

2.SwinTransfomer网络

Swin Transformer是2021年微软研究院发表在ICCV上的一篇文章，并且已经获得ICCV 2021 best paper的荣誉称号。Swin Transformer网络是Transformer模型在视觉领域的又一次碰撞。该论文一经发表就已在多项视觉任务中霸榜。

3.网络效果

测试集上效果：

混淆矩阵：

ROC曲线

4.代码获取

感兴趣的朋友可以关注最后一行

import torch
import torch.nn as nn
import numpy as np
import pandas as pd
from swin_transformer import swin_t
from getCWTImage import CWT
import warnings
from torch import optim
from torch.utils.data import DataLoader
from sklearn.metrics import f1_score, accuracy_score, precision_score, recall_score, roc_curve, auc, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
from pytorch_lightning import seed_everything

warnings.filterwarnings("ignore")
seed_everything(0)
# https://mbd.pub/o/bread/mbd-ZpmTmp9u