基于多分辨率注意力脉冲神经网络的机械振动信号故障诊断算法(Pytorch)
本文提出了一种基于多分辨率注意力脉冲神经网络的轴承故障诊断算法。该方法首先通过滑动窗口对振动信号进行分割,采用3种不同尺度的卷积核并行提取多分辨率特征。算法核心是引入注意力机制的脉冲神经元,模拟生物神经元的膜电位动态过程,将连续信号转换为脉冲序列。通过高效通道注意力模块和残差连接的注意力脉冲基础块进行深层特征提取,最后整合时间步信息实现分类。实验结果表明该算法能有效提取振动信号特征,实现精确故障诊
算法首先对原始的轴承振动信号进行预处理,通过滑动窗口将连续的振动信号分割为固定长度的信号片段。接着,算法采用多分辨率特征提取策略,通过三个并行的卷积分支分别使用3×3、5×5和7×7三种不同尺度的卷积核对信号进行初步特征提取,每个分支都包含卷积层、批量归一化层和注意力脉冲神经元,在脉冲神经元中模拟生物神经元的膜电位积分、脉冲发放和重置过程,将连续信号转换为脉冲序列,随后进行平均池化降维和第二次卷积操作以提取更深层次特征。
三个分支提取的特征经过拼接后,通过高效通道注意力模块对通道特征进行重校准,增强重要特征通道的权重。处理后的特征输入到注意力脉冲神经元中进行脉冲编码和时空动态处理,模拟生物神经网络的信息传递机制。
随后,特征依次通过两个注意力脉冲基础块进行深层特征提取,每个基础块都包含卷积操作、批量归一化、通道注意力、空间注意力和脉冲发放等多个步骤,并通过残差连接保持梯度流动。最后,通过全局平均池化将时空特征压缩为固定维度的特征向量,通过全连接层将特征映射到故障类别空间,得到初步分类结果。由于网络处理了多个时间步的脉冲序列,最终将各个时间步的输出进行平均,整合时间维度上的信息,形成最终分类决策。整个训练过程使用交叉熵损失函数计算预测误差,通过反向传播和替代梯度法更新网络参数,采用学习率衰减策略优化训练过程,实现端到端的轴承故障诊断。

算法步骤
数据准备阶段:将连续的轴承振动信号通过滑动窗口技术分割为固定长度的信号片段,确保每个样本包含完整的振动特征信息。
多尺度特征提取阶段:通过三个并行的卷积神经网络分支,分别使用小、中、大三种不同尺度的卷积核对输入信号进行初步特征提取,每个分支都包含卷积操作、批量归一化处理和注意力脉冲神经元转换。
脉冲编码阶段:在每个分支的注意力脉冲神经元中,模拟生物神经元的动态过程,将连续的特征值转换为离散的脉冲序列,包括膜电位积分、阈值比较、脉冲发放和膜电位重置等步骤。
特征融合阶段:将三个分支提取的多尺度特征进行拼接,形成综合特征表示,然后通过高效通道注意力机制对特征通道进行权重调整,增强重要特征通道的影响。
深层特征提取阶段:将融合后的特征输入到两个串联的注意力脉冲基础块中,每个基础块都包含卷积层、批量归一化层、通道注意力模块、空间注意力模块和脉冲发放层,通过残差连接保持信息流动。
特征压缩阶段:通过全局平均池化操作将时空特征压缩为固定维度的特征向量,保留最重要的特征信息,减少参数数量。
分类决策阶段:将压缩后的特征向量输入到全连接分类器中,映射到故障类别空间,得到每个时间步的分类结果,然后将多个时间步的输出进行平均,形成最终分类决策。
训练优化阶段:使用交叉熵损失函数衡量预测结果与真实标签之间的差异,通过反向传播算法计算梯度,使用替代梯度法解决脉冲函数的不可微问题,更新网络参数,采用学习率衰减策略提高训练稳定性。
训练结果分析:
完美的分类性能:训练集和验证集都达到了100%准确率
损失极小:训练损失和验证损失都降到了0.0000级别
学习率调度有效:第30个epoch学习率从0.01降到了0.001
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
import os
import json
import time
import math
from datetime import datetime
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import signal
from scipy.fft import fft, fftfreq
from sklearn.metrics import confusion_matrix, classification_report
# ============================================================================
# 1. 配置类 (Configuration Class)
# ============================================================================
class Config:
"""
配置类 - 定义所有超参数和路径
Configuration Class - Define all hyperparameters and paths
"""
# 设备配置 Device Configuration
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
# 数据路径 Data Paths
DATA_DIR = './data'
MODEL_SAVE_DIR = './trained_models'
VISUALIZATION_DIR = './visualizations'
# 数据参数 Data Parameters
SEQUENCE_LENGTH = 1024 # 样本长度 Sample length
SAMPLE_RATE = 12000 # 采样频率 Sampling frequency (Hz)
OVERLAP_RATIO = 0.5 # 滑动窗口重叠率 Sliding window overlap ratio
# 训练参数 Training Parameters
BATCH_SIZE = 64
EPOCHS = 100
LEARNING_RATE = 0.01
NUM_CLASSES = 4 # 4种故障类型 4 fault types
# 优化器参数 Optimizer Parameters
MOMENTUM = 0.9
WEIGHT_DECAY = 1e-4
# 学习率调度器参数 Learning Rate Scheduler Parameters
STEP_SIZE = 30
GAMMA = 0.1
# 脉冲神经网络参数 Spiking Neural Network Parameters
TIMESTEP = 4 # SNN时间步数 SNN timesteps
VTH = 1.0 # 脉冲神经元阈值 Spiking neuron threshold
A = 1.0 # 梯度近似参数 Gradient approximation parameter
TAU = 0.5 # 膜电位衰减常数 Membrane potential decay constant
ALPHA = 0.5 # 注意力参数 Attention parameter
# 数据划分比例 Data Split Ratio
TRAIN_RATIO = 0.7
VAL_RATIO = 0.15
TEST_RATIO = 0.15
# 随机种子 Random Seed
SEED = 42
# 故障标签 Fault Labels
FAULT_LABELS = {
0: 'Normal',
1: 'Ball Fault (0.007")',
2: 'Inner Race Fault (0.007")',
3: 'Outer Race Fault (0.007")'
}
# 文件名到标签的映射 File to label mapping
FILE_LABELS = {
'98raw.txt': 0, # 正常状态 Normal condition
'106raw.txt': 1, # 滚珠故障 Ball fault
'119raw.txt': 2, # 内圈故障 Inner race fault
'131raw.txt': 3 # 外圈故障 Outer race fault
}
# ============================================================================
# 2. 脉冲神经网络核心组件 (Spiking Neural Network Core Components)
# ============================================================================
class StepFunction(torch.autograd.Function):
"""
阶跃函数,带自定义梯度
Step function with custom gradient
"""
@staticmethod
def forward(ctx, input):
"""
前向传播
Forward pass
Args:
input: 输入张量 Input tensor
Returns:
阶跃输出 Step output
"""
ctx.save_for_backward(input)
output = torch.ge(input, 0)
return output.float()
@staticmethod
def backward(ctx, grad_output):
"""
反向传播
Backward pass
Args:
grad_output: 输出梯度 Output gradient
Returns:
输入梯度 Input gradient
"""
input, = ctx.saved_tensors
grad_input = grad_output.clone()
# 使用矩形函数近似梯度 Use rectangle function to approximate gradient
hu = (abs(input) < (Config.A / 2)) / Config.A
return grad_input * hu
担任《Mechanical System and Signal Processing》《中国电机工程学报》《宇航学报》《控制与决策》等期刊审稿专家,擅长领域:信号滤波/降噪,机器学习/深度学习,时间序列预分析/预测,设备故障诊断/缺陷检测/异常检测
参考文章:
基于多分辨率注意力脉冲神经网络的机械振动信号故障诊断算法(Pytorch) - 哥廷根数学学派的文章
https://zhuanlan.zhihu.com/p/1995631120000254737
魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。
更多推荐


所有评论(0)