大家好，这里是公众号：淘个代码！公众号改版信息流之后很多读者压根收不到小淘发出的文章，建议大家星标公众号！
由于微信公众平台改变了推送机制，凡是点赞、在看、星标公众号的朋友，都会优先接收到文章推送，所以大家一定要记得星标一下公众号，读完文章后顺手点一下，避免错过文章！
前言
本期整合了10种时频图 + 12种注意力机制，嵌入CNN网络，实现故障诊断。首先将1D振动信号转换为2D时频图，再利用2DCNN + 注意力机制进行故障识别。一共实现了10种时频转换方法和12种注意力机制,可自由组合，共计130种实验方案，方便横向对比不同方法的效果。
一、10种时频图转换方法
将轴承信号分割成2048个数据点的1D振动信号，然后将1D信号转换为64×64的2D图像，实现的方法有：
序号 | 方法 | 说明 |
1 | STFT | 短时傅里叶变换，最经典的时频分析方法 |
2 | CWT | 连续小波变换，多尺度时频分析 |
3 | Mel频谱图 | Mel频率谱，模拟人耳感知特性 |
4 | GASF | Gramian角求和场，保留时间相关性 |
5 | GADF | Gramian角差分场，捕捉信号变化趋势 |
6 | MTF | Markov转移场，编码动态转移概率 |
7 | RP | 递归图，揭示信号的递归结构 |
8 | Hilbert谱 | Hilbert-Huang变换，瞬时频率分布 |
9 | S变换 | Stockwell变换，兼具STFT和CWT优点 |
10 | WVD | Wigner-Ville分布，高时频分辨率 |
转换后的效果如下图所示（以Normal类第1个样本为例）：
每种方法从不同角度揭示信号的时频特性，各有侧重。例如CWT擅长捕捉多尺度特征，GASF/GADF保留了时间序列的内在关联，RP则突出了信号的周期性结构。
二、12种注意力机制
在CNN中嵌入注意力模块，帮助模型聚焦于关键特征区域，实现的注意力机制有：
序号 | 注意力机制 | 类型 | 核心思想 |
1 | SE-Net | 通道 | Squeeze-Excitation，通道加权 |
2 | CBAM | 通道+空间 | 先通道后空间，双重注意力 |
3 | ECA | 通道 | 高效通道注意力，1D卷积替代全连接 |
4 | GAM | 全局 | 通道-空间联合建模 + Channel Shuffle |
5 | CA | 坐标 | 水平+垂直方向位置编码 |
6 | Triplet | 跨维度 | 三分支交叉维度交互 |
7 | SimAM | 无参数 | 基于神经科学的自相似度估计 |
8 | BAM | 瓶颈 | 通道门+空间门并行融合 |
9 | SK-Net | 多尺度 | 多尺度卷积核自适应选择 |
10 | GC-Net | 全局上下文 | 全局上下文信息建模 |
11 | A2-Net | 双重 | 聚合-分配双重注意力 |
12 | DANet | 位置+通道 | 位置注意力与通道注意力并行 |
三、2DCNN网络结构
基础网络为4层卷积 + 全连接分类器，注意力模块插在最后一层卷积之后：
输入 (batch, 3, 64, 64) ↓ Conv2d(3→32, 3×3) → BN → ReLU → MaxPool(2×2) ↓ Conv2d(32→64, 3×3) → BN → ReLU → MaxPool(2×2) ↓ Conv2d(64→128, 3×3) → BN → ReLU → MaxPool(2×2) ↓ Conv2d(128→256, 3×3) → BN → ReLU → AdaptiveAvgPool(4×4) ↓ 【注意力机制插入点】 ↓ Flatten → FC(4096→256) → Dropout(0.3) → ReLU → FC(256→10) ↓ 输出 10类故障分类四、数据集
采用CWRU（凯斯西储大学）轴承数据集，包含10种工况。
类别 | 说明 |
Normal | 正常 |
IR007 | 内圈故障 0.007英寸 |
B007 | 滚动体故障 0.007英寸 |
OR007@6 | 外圈故障 0.007英寸（6点钟方向） |
IR014 | 内圈故障 0.014英寸 |
B014 | 滚动体故障 0.014英寸 |
OR014@6 | 外圈故障 0.014英寸（6点钟方向） |
IR021 | 内圈故障 0.021英寸 |
B021 | 滚动体故障 0.021英寸 |
OR021@6 | 外圈故障 0.021英寸（6点钟方向） |
每类320个样本，每个样本长度为2048个采样点，并添加了高斯白噪声以增加诊断难度。专门写了一个脚本可以自由更改噪声强度。经时频变换后转为(64, 64, 3)的2D图像。
五、结果展示
采用STFT变换，融合BAM机制的CNN结果如下：
STFT_BAM-CNN2D：采用ST变换，融合SK机制的CNN结果如下：
一共有130种组合，就不一一展示了。
除此之外还设计了一个脚本专门用于结果对比：
环境要求
• Python 3.10+
• PyTorch 2.0+（支持CUDA）
• 其他依赖：
pip install numpy scipy scikit-learn scikit-image pandas matplotlib seaborn Pillow PyWavelets项目结构
├── signal_to_2d.py # 10种1D→2D时频变换方法 ├── convert_dataset.py # 批量数据转换脚本（生成npy+png）├── main_2d.py # 训练主脚本├── plot_def/│ └── plot_acc_loss_2d.py # 多模型/多方法对比绘图├── model/│ ├── CNN2D_AM.py # 13种2D CNN模型定义│ └── attention2d/ # 12种2D注意力机制模块│ ├── se2d.py # SE-Net│ ├── cbam2d.py # CBAM│ ├── eca2d.py # ECA│ ├── gam2d.py # GAM│ ├── ca2d.py # Coordinate Attention│ ├── triplet2d.py # Triplet Attention│ ├── simam2d.py # SimAM│ ├── bam2d.py # BAM│ ├── sk2d.py # SK-Net│ ├── gc2d.py # GC-Net│ ├── a2net2d.py # A2-Net│ └── danet2d.py # DANet├── ulit/│ ├── DataLoad2D.py # 2D数据集加载│ ├── CustomDataset2D.py # PyTorch Dataset│ ├── toolFunc2D.py # 模型选择工具│ ├── acc.py # 准确率计算│ └── init_seed.py # 随机种子初始化├── run_all.py # 一键批量训练所有组合├── dataset/│ ├── Add_Noise_CWRU/ # 原始1D数据（.mat文件）│ ├── Add_Noise_CWRU_2D/ # 转换后2D数据（.npy文件）│ └── Add_Noise_CWRU_2D_img/ # 转换后的可视化图片（.png文件） └── result_2d/ # 训练结果输出目录运行步骤
第一步：数据转换
将原始1D振动信号（.mat文件）转换为2D时频图（.npy+.png）。
转换单种方法：
python convert_dataset.py --method STFT一次性转换全部10种方法：
python convert_dataset.py --method all可选的10种时频变换方法：
方法 | 说明 |
STFT | 短时傅里叶变换 |
CWT | 连续小波变换 |
Mel | Mel频谱图 |
GASF | Gramian角求和场 |
GADF | Gramian角差分场 |
MTF | Markov转移场 |
RP | 递归图 |
Hilbert | Hilbert谱 |
ST | S变换（Stockwell） |
WVD | Wigner-Ville分布 |
转换完成后：
•
.npy数据保存在dataset/Add_Noise_CWRU_2D/{方法名}/目录下•
.png可视化图片保存在dataset/Add_Noise_CWRU_2D_img/{方法名}/目录下
第二步：模型训练
python main_2d.py --tf_method STFT --use_model CNN2D --epochs 30参数说明：
参数 | 默认值 | 说明 |
--tf_method | STFT | 时频变换方法（需先完成对应数据转换） |
--use_model | CNN2D | 模型名称（见下表） |
--epochs | 30 | 训练轮数 |
--batch-size | 32 | 批次大小 |
--lr | 0.001 | 学习率 |
--data | ./dataset | 数据集根目录 |
--save_dir | ./result_2d | 结果保存目录 |
可选的13种模型：
模型名称 | 注意力机制 |
CNN2D | 无（基准模型） |
SE-CNN2D | SE-Net |
CBAM-CNN2D | CBAM |
ECA-CNN2D | ECA |
GAM-CNN2D | GAM |
CA-CNN2D | Coordinate Attention |
Triplet-CNN2D | Triplet Attention |
SimAM-CNN2D | SimAM |
BAM-CNN2D | BAM |
SK-CNN2D | SK-Net |
GC-CNN2D | GC-Net |
A2-CNN2D | A2-Net |
DANet-CNN2D | DANet |
训练示例：
# 使用STFT + CBAM注意力机制 python main_2d.py --tf_method STFT --use_model CBAM-CNN2D --epochs 30 # 使用CWT + SE注意力机制 python main_2d.py --tf_method CWT --use_model SE-CNN2D --epochs 30 # 使用Mel频谱图 + 基准CNN（无注意力） python main_2d.py --tf_method Mel --use_model CNN2D --epochs 30一键批量训练所有组合（10种时频 × 13种模型 = 130个组合）：
python run_all.py也可以自定义参数：
# 自定义训练轮数 python run_all.py --epochs 50 # 只跑指定的时频方法（多个用空格分隔） python run_all.py --tf_method STFT CWT Mel # 只跑指定的模型 python run_all.py --use_model CNN2D CBAM-CNN2D SE-CNN2D # 组合指定（例如：2种时频 × 3种模型 = 6个组合） python run_all.py --tf_method STFT CWT --use_model CNN2D CBAM-CNN2D SE-CNN2D脚本会依次运行每个组合，实时显示进度[1/130]，训练结束后汇总成功/失败数量。
第三步：查看结果
训练完成后，结果保存在result_2d/ADD_Noise_CWRU/{时频方法}_{模型名}/目录下：
result_2d/ADD_Noise_CWRU/STFT_CBAM-CNN2D/ ├── STFT_CBAM-CNN2D_acc_loss.png # 准确率和损失曲线 ├── STFT_CBAM-CNN2D_混淆矩阵.png # 混淆矩阵热图 ├── STFT_CBAM-CNN2D_测试集原始特征分布.png # t-SNE原始特征可视化 ├── STFT_CBAM-CNN2D_全连接层特征分布.png # t-SNE全连接层特征可视化 ├── STFT_CBAM-CNN2D_train_test_result.csv # 逐epoch训练/测试记录 ├── confusion_matrix.txt # 混淆矩阵文本 └── model_30.pth # 训练好的模型权重第四步：对比绘图（可选）
训练多个模型后，可以使用对比绘图脚本进行结果对比。
对比同一时频方法下的不同注意力机制：
python plot_def/plot_acc_loss_2d.py --mode models --tf_method STFT对比同一模型下的不同时频方法：
python plot_def/plot_acc_loss_2d.py --mode tf_methods --model CNN2D全部对比（需要先完成所有组合的训练）：
python plot_def/plot_acc_loss_2d.py --mode all总结
本项目提供了一套完整、灵活、可扩展的2D CNN轴承故障诊断框架：
•10种时频转换覆盖了主流的信号-图像编码方法
•12种注意力机制涵盖通道、空间、坐标、无参数等多种类型
•130种组合方案可一键训练、一键对比
• 自动保存准确率曲线、混淆矩阵、t-SNE特征分布等完整结果
拿到数据后，三行命令就能跑通全流程，非常适合作为时频分析 + 注意力机制方向的研究基线或课题参考。
欢迎关注公众号：淘个代码 代码获取跳转链接：
https://mbd.pub/o/bread/YZWbmZpxaQ==
欢迎关注公众号：淘个代码
