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

        Swintransformer是microsoft在2021年提出的方法（原版code链接），本文通过将其与小波时频图结合起来，共同用于轴承故障诊断中，目前还没有人将这个方法用于故障诊断哟。

1.Swintransformer结构

        上图展示的是官方的模型结构图，实际上是tiny模型，。它首先通过利用patch partition将输入的图像分割为非重叠 patch。其次采用linear embedding，将patch投影到维度C，然后交替使用窗口注意力机制与移位窗口注意力，并采用patch merging进行下采样操作（起到CNN中池化层的作用）。与基于resnet的结构类似，可以得到不同stage的不同下采样倍数的特征图，用于各种下游任务，比如可以取stage2-4的特征图输入FPN（如yolov4的结构），然后接yolohead进行目标检测等。 

2.小波时频图

        传统轴承故障诊断是采用各种特征提取方法对一维轴承信号进行特征提取，如均值方差，HHT包络谱，FFT频谱，小波能量谱等，变换后的特征依旧是一维数据。本文采用小波时频图技术将一维轴承信号转换为2维（3通道真彩图）的时频图，转换后的数据可以输入CNN、MLP-Mixer、visiontransformer等中进行图像分类，从而实现故障诊断。

3.实例验证

代码列表

       1）数据准备

       采用西储大学轴承故障诊断数据集，48K/0HP数据，共10类故障（正常作为一类特殊的故障类型），划分后每个样本的采样点为1024，每类故障各100个样本，因此一共1000个样本，然后7:2:1划分训练集，验证集，与测试集。

      2）小波时频图

        3）SwinTransformer故障诊断分类

        参考代码地址如下SwinTransformer细节及代码实现（pytorch版本），将其与时频图结合起来，其中输入大小是64x64x3，patchsize是4，类别数为10。经过100次训练的正确率曲线与损失曲线如下图所示。

        将测试集输入3）中训练好的模型，得到正确率为99%。

        4）特征可视化

        提取分类层前一层的输出，也就是判别特征，采用tsne进行可视化。

5 代码

代码见评论区，另外链接中还有更多博客及代码。