MACNN轴承故障诊断模型代码功能全解析

基于注意力的多尺度卷积神经网络轴承故障诊断 针对传统方法在噪声环境下诊断精度低的问题提出了一种多尺度卷积神经网络的滚动轴承故障诊断方法 首先构建多尺度卷积提取不同尺度的故障特征同时引入通道注意力自适应地选择包含故障特征的通道来提高模型的抗噪能力抑制噪声干扰此外利用自适应大小的一维卷积调整不同尺度的特征通道权重自适应融合不同尺度的特征提高判别性特征提取能力最后通过凯斯西储大学开源滚动轴承数据集CWRU进行验证证明了所提方法对有效性 参考文献2023年吉林大学学报EI《基于注意力的多尺度卷积神经网络轴承故障诊断》 ●数据预处理支持1维原始数据 ●网络模型1DMACNN ●数据集凯斯西储大学开源滚动轴承数据集CWRU、十分类 ●网络框架pytorch ●结果输出损失曲线图、准确率曲线图、混淆矩阵、tsne图 ●准确率测试集100% ●使用对象初学者 ●代码保证代码注释详细、即拿即可跑通一、项目核心定位本项目是基于PyTorch框架开发的轴承故障诊断系统核心为多尺度注意力卷积神经网络MACNN。该系统针对传统卷积神经网络在噪声环境下轴承故障诊断精度低的问题通过融合多尺度特征提取与通道注意力机制实现了对滚动轴承10类故障含正常状态的高精度识别。项目提供了从数据集加载、模型训练、性能评估到特征可视化的完整流程支持CWRU轴承数据集可直接用于噪声干扰、负载变化等复杂工况下的轴承故障诊断任务。二、代码架构总览项目文件按功能模块划分结构清晰共包含13个核心文件具体组织如下MACNN_Project/ ├── main.py # 核心主程序整合训练、测试、结果保存全流程 ├── model/ # 模型定义模块存放网络结构及注意力机制实现 │ ├── MACNN.py # MACNN主模型架构定义 │ ├── SEnet.py # SE通道注意力机制实现 │ ├── ECAnet.py # ECA高效通道注意力机制实现 │ └── __init__.py # 模块初始化空文件 ├── ulit/ # 工具辅助模块提供数据处理、指标计算等基础功能 │ ├── CWRU.py # CWRU数据集加载与划分 │ ├── cwru_datasets.py # 自定义数据集类适配PyTorch DataLoader │ ├── acc.py # 准确率与损失值统计工具 │ ├── init_seed.py # 随机种子初始化保证实验可复现 │ └── __init__.py # 模块初始化空文件 └── plot_def/ # 可视化模块生成训练曲线与特征可视化图 ├── plot_acc_loss.py # 训练/测试准确率、损失曲线绘制 ├── t_SNE_feat_save.py # t-SNE特征降维可视化 └── __init__.py # 模块初始化空文件三、核心文件功能深度解析3.1 主程序文件main.py作为项目入口负责串联数据加载、模型训练、测试评估、结果保存等全流程支持命令行参数配置是整个系统的中枢。3.1.1 关键功能模块环境初始化- 自动检测GPU/CPU环境优先使用CUDA加速无GPU时自动切换至CPU。- 随机种子初始化调用init_seed()函数固定Python、NumPy、PyTorch的随机种子默认123确保实验结果可复现。- 命令行参数解析支持配置数据集路径、训练轮数、批次大小、学习率、权重衰减等关键参数灵活适配不同实验需求。模型与优化器配置- 模型加载根据--use_model参数加载MACNN模型并将模型移至指定计算设备GPU/CPU。- 损失函数采用交叉熵损失函数nn.CrossEntropyLoss适配多分类任务。- 优化器使用Adam优化器学习率默认0.001权重衰减1e-4抑制过拟合。- 学习率调度采用StepLR调度器每10轮--stepsize学习率衰减为原来的0.1倍--gamma。数据处理流程- 数据集加载通过CWRU类读取指定路径下的CWRU数据集自动识别故障类别标签。- 数据划分按8:2比例划分训练集与测试集采用分层抽样stratify保证各类别样本分布均衡。- 批量加载通过torch.utils.data.DataLoader实现数据批量读取支持多线程加载默认0线程提升训练效率。训练与测试循环- 训练过程迭代训练指定轮数默认50轮每轮遍历训练集完成前向传播、损失计算、反向传播与参数更新实时统计训练准确率与损失值。- 测试过程每轮训练后在测试集上评估模型性能统计测试准确率、损失值记录真实标签与预测标签。- 结果保存训练结束后默认第49轮保存模型权重文件.pth、混淆矩阵confusion_matrix.txt、分类报告同时将每轮的准确率与损失值写入CSV文件。可视化输出- 绘制训练/测试准确率曲线与损失曲线直观展示模型收敛过程。- 生成混淆矩阵热力图清晰呈现各类故障的分类效果与混淆情况。3.1.2 核心代码片段解析def train(train_loader, model, criterion, optimizer, epoch, lr_scheduler, device): losses AverageMeter(Loss, :.4f) train_acc AverageMeter(Acc, :.4f) model.train() # 开启训练模式启用Dropout、BN层更新 for i, (data, label) in enumerate(train_loader): model.zero_grad() optimizer.zero_grad() input data.to(device) label label.to(device) output model(input) # 前向传播输入数据经过模型得到预测输出 loss criterion(output, label.long()) # 计算预测值与真实值的损失 losses.update(loss.item(), label.size(0)) # 更新损失统计 # 计算训练准确率 _, predicted torch.max(output, 1) # 获取预测概率最大的类别 accuracy (predicted label).sum().item() / label.size(0) train_acc.update(accuracy, label.size(0)) # 更新准确率统计 loss.backward() # 反向传播计算梯度 optimizer.step() # 参数更新根据梯度调整模型权重 lr_scheduler.step() # 学习率更新 print(fEpoch:[{epoch}] train_Acc:{train_acc.avg:.4f} train_Loss:{losses.avg:.4f}) return train_acc.avg, losses.avg该函数实现了单轮训练的核心逻辑通过AverageMeter类实时统计批次损失与准确率model.train()开启训练模式确保BN层和Dropout层正常工作反向传播过程通过loss.backward()计算梯度optimizer.step()完成参数更新最终返回本轮平均准确率与损失值。3.2 模型定义模块model文件夹该模块包含MACNN主模型及两种注意力机制的实现是整个系统的核心算法载体。3.2.1 MACNN.py主模型架构MACNN模型通过多尺度卷积提取特征、注意力机制强化有效信息、自适应融合特征实现高抗噪性故障诊断其结构严格遵循论文设计。层结构定义- 宽卷积层conv1 nn.Conv1d(1, 32, kernelsize64, stride2, padding1)输入为1维振动信号输出32通道特征大卷积核可有效抑制噪声提取长时尺度特征。- 批量归一化BN层bn nn.BatchNorm1d(32)加速模型收敛缓解梯度消失问题。- 最大池化层maxpool1 nn.MaxPool1d(kernelsize2, stride2)降维并保留关键特征减少计算量。- 多尺度卷积块包含两组多尺度卷积层每组由3个不同核尺寸5×1、7×1、9×1的卷积层并行组成分别提取不同频率范围的故障特征适配不同类型、不同直径的轴承故障。- 注意力层se1 SEBlock(64)、se2 SEBlock(128)分别作用于两组多尺度卷积的输出自适应调整通道权重eca ECABlock(1283)作用于融合后的特征进一步强化有效特征。- 全局平均池化层globalavgpooling nn.AdaptiveAvgPool1d(1)将特征图压缩为向量减少参数数量提升模型泛化能力。- 全连接层classificationlayer nn.Linear(1283, self.num_classes)将融合后的特征映射为10类故障的预测概率。前向传播逻辑def forward(self, x): # 宽卷积层特征提取 x F.relu(self.bn(self.conv1(x))) x self.max_pool1(x) # 第一条多尺度支路核尺寸5×1 x1 F.relu(self.conv2_1(x)) x1 self.max_pool2_1(x1) x1 self.se1(x1) x1 F.relu(self.conv3_1(x1)) x1 self.max_pool3_1(x1) x1 self.se2(x1) # 第二条多尺度支路核尺寸7×1 x2 F.relu(self.conv2_2(x)) x2 self.max_pool2_2(x2) x2 self.se1(x2) x2 F.relu(self.conv3_2(x2)) x2 self.max_pool3_2(x2) x2 self.se2(x2) # 第三条多尺度支路核尺寸9×1 x3 F.relu(self.conv2_3(x)) x3 self.max_pool2_3(x3) x3 self.se1(x3) x3 F.relu(self.conv3_3(x3)) x3 self.max_pool3_3(x3) x3 self.se2(x3) # 多尺度特征融合 ECA注意力加权 x torch.cat([x1, x2, x3], dim1) x self.eca(x) # 特征压缩与分类 x self.global_avg_pooling(x) x x.view(x.size(0), -1) # 展平特征向量 x self.classification_layer(x) return x前向传播过程分为三个关键阶段首先通过宽卷积层完成初步特征提取与噪声抑制然后三条多尺度支路并行提取特征经SE注意力机制强化有效通道最后拼接多尺度特征通过ECA注意力机制自适应调整融合特征权重最终经全连接层输出分类结果。3.2.2 SEnet.pySE通道注意力机制基于Squeeze-and-ExcitationSE模块通过显式建模通道间的依赖关系自适应校准特征响应提升模型对关键特征的关注度。class SEBlock(nn.Module): def __init__(self, channel, reduction16): super(SEBlock, self).__init__() self.avg_pool nn.AdaptiveAvgPool1d(1) # Squeeze全局平均池化 self.fc nn.Sequential( # Excitation全连接层学习通道权重 nn.Linear(channel, channel // reduction, biasFalse), nn.ReLU(inplaceTrue), nn.Linear(channel // reduction, channel, biasFalse), nn.Sigmoid() ) def forward(self, x): b, c, _ x.size() # Squeeze压缩空间维度得到全局通道信息 y self.avg_pool(x).view(b, -1) # Excitation学习通道权重 y self.fc(y).view(b, c, 1) # 特征加权将权重应用于原始特征 return x * y.expand_as(x)SE模块工作流程首先通过全局平均池化Squeeze将每个通道的特征图压缩为单个数值表征该通道的全局信息然后通过两层全连接层Excitation学习通道间的依赖关系输出每个通道的权重最后将权重与原始特征相乘实现对关键通道的强化和无用通道的抑制。3.2.3 ECAnet.pyECA高效通道注意力机制优化SE模块的计算效率通过1D卷积替代全连接层在保持性能的同时减少参数数量适用于特征融合后的通道权重调整。class ECABlock(nn.Module): def __init__(self, channel, gamma2, b1): super(ECABlock, self).__init__() self.avg_pool nn.AdaptiveAvgPool1d(1) # 1D卷积学习通道权重替代SE的全连接层 self.conv nn.Sequential( nn.Conv1d(channel, channel // gamma, kernel_size1), nn.ReLU(inplaceTrue), nn.Conv1d(channel // gamma, channel, kernel_size1), nn.Sigmoid() ) self.b b # 偏置项增强模型鲁棒性 def forward(self, x): y self.avg_pool(x) # 调整维度以适配1D卷积 y self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1) # 特征加权引入偏置项避免权重过拟合 return x * (self.b y)ECA模块的核心改进使用1D卷积替代SE模块中的全连接层减少了参数计算量引入偏置项b避免权重过度依赖输入特征提升模型在噪声环境下的鲁棒性。3.3 工具辅助模块ulit文件夹提供数据集处理、性能指标计算、随机种子初始化等基础功能为核心流程提供支撑。3.3.1 CWRU.py数据集加载与划分专门用于处理CWRU轴承数据集实现数据路径读取、标签分配与训练/测试集划分。class CWRU(object): def __init__(self, root_dir, test_size0.2, transformNone): self.root_dir root_dir # 数据集根路径 self.test_size test_size # 测试集比例 self.data_pd self.load_cwru_data() # 加载数据路径与标签 self.transform transform def load_cwru_data(self): data {data: [], label: []} # 遍历数据集目录分配标签按文件夹名称排序 for class_label, class_name in enumerate(os.listdir(self.root_dir)): class_path os.path.join(self.root_dir, class_name) if os.path.isdir(class_path): # 读取该类别下所有.mat文件路径 for file_name in os.listdir(class_path): file_path os.path.join(class_path, file_name) data[data].append(file_path) data[label].append(class_label) return pd.DataFrame(data) def train_test_split_order(self): # 分层抽样划分训练集与测试集保证标签分布均衡 train_pd, test_pd, _, _ train_test_split( self.data_pd, self.data_pd[label], test_sizeself.test_size, stratifyself.data_pd[label], random_state123 ) # 构建自定义数据集实例 train_dataset CustomCWRUDataset(train_pd, self.transform) test_dataset CustomCWRUDataset(test_pd, self.transform) return train_dataset, test_dataset核心功能loadcwrudata()方法遍历数据集目录将每个文件夹视为一个故障类别分配对应的标签并存储文件路径与标签到DataFrametraintestsplit_order()方法使用分层抽样划分数据集确保训练集与测试集中各类别样本比例一致避免数据偏斜影响模型性能。3.3.2 cwru_datasets.py自定义数据集类实现PyTorch的Dataset接口用于读取单个.mat格式的数据文件将其转换为模型可处理的Tensor格式。class CustomCWRUDataset(Dataset): def __init__(self, data_pd, transformNone): self.data_pd data_pd # 包含文件路径与标签的DataFrame self.transform transform def __len__(self): return len(self.data_pd) # 返回数据集样本总数 def __getitem__(self, idx): # 获取当前样本的文件路径与标签 file_path self.data_pd.iloc[idx][data] label int(self.data_pd.iloc[idx][label]) # 读取.mat文件中的振动信号sample字段存储信号数据 data loadmat(file_path)[sample].transpose() # 转换为Tensor格式float类型 data torch.tensor(data).float() # 应用数据变换如归一化、增强等可选 if self.transform: data self.transform(data) return data, label # 返回数据标签对该类是PyTorch数据加载的核心组件getitem方法实现了单个样本的读取逻辑读取.mat文件中的振动信号转置为模型要求的维度转换为Tensor格式后返回支持自定义数据变换如归一化、数据增强等。3.3.3 acc.py性能指标统计工具AverageMeter类用于实时统计训练/测试过程中的准确率、损失值等指标支持批次更新与平均值计算。class AverageMeter(object): Computes and stores the average and current value def __init__(self, name, fmt:f): self.name name # 指标名称如Loss、Acc self.fmt fmt # 格式化字符串 self.reset() # 初始化统计参数 def reset(self): self.val 0 # 当前批次的指标值 self.avg 0 # 累计平均指标值 self.sum 0 # 累计指标总和 self.count 0 # 累计样本数 def update(self, val, n1): self.val val # 更新当前批次值 self.sum val * n # 累计总和乘以样本数支持批次大小不一致 self.count n # 累计样本数 self.avg self.sum / self.count # 计算平均值核心作用在训练/测试过程中每处理一个批次就调用update()方法更新统计信息最终通过avg属性获取整个数据集的平均指标值方便实时监控模型性能。3.3.4 init_seed.py随机种子初始化固定所有随机数生成器的种子确保实验结果可复现避免因随机因素导致的结果波动。def init_seed(seed123): seed seed np.random.seed(seed) # NumPy随机种子 random.seed(seed) # Python随机种子 torch.manual_seed(seed) # PyTorch CPU随机种子 torch.cuda.manual_seed(seed) # PyTorch GPU随机种子 torch.cuda.manual_seed_all(seed) # 多GPU场景下的随机种子 cudnn.benchmark False # 禁用CUDA的非确定性算法 cudnn.deterministic True # 启用CUDA的确定性算法关键作用在深度学习实验中随机种子的固定是保证结果可复现的关键。该函数覆盖了Python、NumPy、PyTorchCPU/GPU的所有随机数生成器同时配置CUDA的确定性算法确保每次运行代码的结果一致。3.4 可视化模块plot_def文件夹提供训练过程可视化与特征可视化功能帮助用户直观分析模型性能与特征提取效果。3.4.1 plot_acc_loss.py训练曲线绘制读取main.py保存的CSV文件绘制训练/测试准确率曲线与损失曲线支持多模型对比当前默认MACNN。基于注意力的多尺度卷积神经网络轴承故障诊断 针对传统方法在噪声环境下诊断精度低的问题提出了一种多尺度卷积神经网络的滚动轴承故障诊断方法 首先构建多尺度卷积提取不同尺度的故障特征同时引入通道注意力自适应地选择包含故障特征的通道来提高模型的抗噪能力抑制噪声干扰此外利用自适应大小的一维卷积调整不同尺度的特征通道权重自适应融合不同尺度的特征提高判别性特征提取能力最后通过凯斯西储大学开源滚动轴承数据集CWRU进行验证证明了所提方法对有效性 参考文献2023年吉林大学学报EI《基于注意力的多尺度卷积神经网络轴承故障诊断》 ●数据预处理支持1维原始数据 ●网络模型1DMACNN ●数据集凯斯西储大学开源滚动轴承数据集CWRU、十分类 ●网络框架pytorch ●结果输出损失曲线图、准确率曲线图、混淆矩阵、tsne图 ●准确率测试集100% ●使用对象初学者 ●代码保证代码注释详细、即拿即可跑通核心功能读取CSV文件中的每轮准确率与损失数据。为不同模型配置不同的线型与颜色绘制对比曲线。保存准确率对比图与损失对比图到指定目录方便实验结果分析。3.4.2 t_SNE_feat_save.py特征可视化使用t-SNEt分布随机邻域嵌入算法将高维特征降维至2D空间可视化模型的特征提取效果。核心功能加载训练好的模型权重提取测试集的原始特征与全连接层特征。使用t-SNE算法对高维特征进行降维将其映射到2D平面。绘制不同类别的特征散点图直观展示模型是否能有效区分各类故障特征。保存特征分布图片与特征数据文件用于后续分析。四、数据集规范与要求项目默认使用CWRU凯斯西储大学轴承数据集该数据集是轴承故障诊断领域的常用基准数据集具体规范如下4.1 数据格式文件类型.mat格式MATLAB数据文件。信号存储每个文件的sample字段存储一段轴承振动信号采样频率为12kHz。数据维度原始信号为1D数组经cwru_datasets.py处理后转换为形状为1, N的TensorN为信号长度。4.2 故障类别数据集包含10类故障含正常状态标签与故障类型的对应关系如下标签故障类型故障直径mm说明0Normal-正常轴承状态1IR0070.1778内圈故障2B0070.1778滚动体故障3OR00760.1778外圈故障6点钟位置4IR0140.3556内圈故障5B0140.3556滚动体故障6OR01460.3556外圈故障6点钟位置7IR0210.5534内圈故障8B0210.5534滚动体故障9OR02160.5534外圈故障6点钟位置4.3 目录结构数据集需按以下目录结构组织确保CWRU.py能正确读取文件与分配标签dataset/CWRU/ ├── Normal/ # 正常状态数据文件夹 │ ├── 1.mat │ ├── 2.mat │ └── ... ├── IR007/ # 内圈故障0.1778mm文件夹 │ └── ... ├── B007/ # 滚动体故障0.1778mm文件夹 │ └── ... └── 其他类别文件夹...五、运行流程与环境要求5.1 环境依赖Python版本3.7及以上。PyTorch版本1.10及以上需匹配CUDA版本支持GPU加速。其他依赖包numpy、pandas、matplotlib、seaborn、scikit-learn、scipy、torchvision。5.2 环境配置步骤解压环境压缩包将pytorch1.10.rar解压至无中文路径如G:\fish\evs\pytorch1.10。配置Python解释器- 打开PyCharm进入File → Settings → Project → Python Interpreter。- 点击Add选择System Interpreter浏览至解压路径下的python.exe如G:\fish\evs\pytorch1.10\pytorch1.10\python.exe。- 点击Apply → OK完成解释器配置。5.3 运行步骤准备数据集按上述目录结构组织CWRU数据集确保路径正确。配置参数可通过命令行参数修改训练配置如学习率、训练轮数等或使用默认参数。启动训练运行main.py程序自动开始训练与测试实时输出训练日志。查看结果训练完成后在result/MACNN目录下查看模型权重、混淆矩阵、训练曲线等结果文件。特征可视化运行plotdef/tSNEfeatsave.py生成特征分布可视化图。5.4 命令行参数示例# 使用默认参数训练 python main.py # 自定义学习率、训练轮数与批次大小 python main.py --lr 0.0005 --epochs 100 --batch-size 32 # 更换数据集路径 python main.py --data D:\datasets\CWRU # 禁用模型保存 python main.py --save_model False六、核心性能指标根据论文实验结果该模型在CWRU数据集上的核心性能指标如下噪声环境适应性在信噪比SNR低至-7dB的极端噪声环境下诊断精度仍保持93%以上。跨负载泛化能力在0hp、1hp、2hp、3hp不同负载工况下平均诊断精度达91.91%。分类准确率在无噪声环境下诊断精度接近100%各类故障分类效果优异。训练效率在GPUNVIDIA GeForce 1050Ti环境下50轮训练耗时约15-20分钟效率较高。七、扩展与优化建议数据集扩展支持XJTU-SY等其他轴承数据集只需修改CWRU.py中的数据加载逻辑适配新数据集的文件格式与标签体系。模型优化- 调整多尺度卷积核尺寸如增加11×1卷积核适配更多类型的故障特征。- 引入数据增强技术如添加高斯噪声、信号平移等提升模型泛化能力。- 尝试其他优化器如SGD、RMSprop或学习率调度器如CosineAnnealingLR优化训练过程。功能扩展- 增加实时诊断功能支持读取传感器实时数据并输出故障诊断结果。- 实现模型量化与部署适配嵌入式设备或工业控制系统。- 增加多模型对比功能支持同时训练多个模型并生成对比报告。性能优化- 启用多线程数据加载修改--workers参数提升训练速度。- 使用混合精度训练进一步加速训练过程减少显存占用。八、常见问题排查环境配置错误- 报错“Python executable not found”检查解压路径是否含中文重新配置解释器路径。- 缺少依赖包使用pip install 包名安装缺失的依赖或创建requirements.txt批量安装。数据读取错误- 报错“KeyError: sample”检查.mat文件中是否存在sample字段或修改cwru_datasets.py中的字段名。- 标签错乱确保数据集文件夹名称与故障类别对应且文件夹排序正确。训练效果不佳- 准确率低且波动大增大批次大小、降低学习率、增加训练轮数或检查数据集是否存在标签错误。- 过拟合训练准确率高测试准确率低增加权重衰减、添加Dropout层、扩大数据集或使用数据增强。GPU无法使用- 检查CUDA版本与PyTorch版本是否匹配确保安装了支持CUDA的PyTorch版本。- 运行torch.cuda.is_available()验证GPU可用性若返回False则使用CPU训练。九、总结该项目基于MACNN模型实现了高精度、高抗噪的轴承故障诊断系统代码结构清晰、功能完整从数据加载、模型训练到结果可视化形成了闭环。通过多尺度特征提取与通道注意力机制的融合有效解决了传统CNN在噪声环境下诊断精度低的问题适用于工业场景中复杂工况下的轴承故障诊断。项目支持灵活的参数配置与扩展既可为学术研究提供稳定的实验平台也可通过简单优化部署到实际工业系统中具有较高的实用性与扩展性。