1. 项目背景与核心价值
滚动轴承作为旋转机械的核心部件,其健康状态直接影响设备运行安全。传统故障诊断方法依赖专家经验和信号处理技术,存在人工特征提取复杂、泛化能力有限等问题。基于深度学习的智能诊断方法通过数据驱动方式自动学习特征表示,正在成为工业预测性维护的新范式。
本项目复现的是一种融合注意力机制与1D-CNN的混合模型架构,其创新点主要体现在:
- 采用1D-CNN处理原始振动信号,避免传统时频分析中的信息损失
- 引入注意力模块动态强化关键故障特征
- 在CWRU轴承数据集上实现98.7%的测试准确率
注意:完整复现需要约8GB显存的GPU设备,若使用Colab需升级到Pro版本以获得足够计算资源
2. 环境配置与数据准备
2.1 开发环境搭建
推荐使用conda创建隔离环境:
conda create -n bearing_diagnosis python=3.8 conda activate bearing_diagnosis pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install scikit-learn pandas matplotlib tqdm2.2 数据集处理
使用Case Western Reserve University轴承数据集:
- 下载12k驱动端轴承数据(包含正常状态和7种故障类型)
- 按以下规则构建样本:
def create_samples(signal, window_size=2048, stride=512): samples = [] for i in range(0, len(signal)-window_size, stride): samples.append(signal[i:i+window_size]) return np.array(samples)- 划分训练/验证/测试集(6:2:2比例)
3. 模型架构详解
3.1 1D-CNN特征提取层
class FeatureExtractor(nn.Module): def __init__(self): super().__init__() self.conv_blocks = nn.Sequential( nn.Conv1d(1, 64, kernel_size=64, stride=8, padding=28), nn.BatchNorm1d(64), nn.ReLU(), nn.MaxPool1d(2, stride=2), nn.Conv1d(64, 128, kernel_size=32, stride=4, padding=14), nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(2, stride=2), nn.Conv1d(128, 256, kernel_size=16, stride=2, padding=8), nn.BatchNorm1d(256), nn.ReLU(), nn.MaxPool1d(2, stride=2) )3.2 通道注意力模块实现
class ChannelAttention(nn.Module): def __init__(self, channels, reduction=16): super().__init__() self.avg_pool = nn.AdaptiveAvgPool1d(1) self.max_pool = nn.AdaptiveMaxPool1d(1) self.fc = nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(), nn.Linear(channels // reduction, channels), nn.Sigmoid() ) def forward(self, x): avg_out = self.fc(self.avg_pool(x).squeeze(-1)) max_out = self.fc(self.max_pool(x).squeeze(-1)) return (avg_out + max_out).unsqueeze(-1)4. 训练优化技巧
4.1 损失函数选择
采用标签平滑的CrossEntropyLoss:
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)4.2 学习率调度策略
使用Warmup+Cosine衰减:
scheduler = torch.optim.lr_scheduler.SequentialLR( optimizer, [ torch.optim.lr_scheduler.LinearLR( optimizer, start_factor=0.01, total_iters=5), torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=95) ], milestones=[5] )5. 常见问题排查
5.1 梯度消失问题
现象:验证集准确率长期停滞在50%左右 解决方案:
- 检查BN层是否在训练模式(model.train())
- 适当减小第一层卷积的stride(从8调整为4)
- 添加梯度裁剪:
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
5.2 过拟合处理
当训练/验证准确率差距>15%时:
- 在卷积后添加Dropout层(p=0.2)
- 使用Mixup数据增强:
def mixup_data(x, y, alpha=0.4): lam = np.random.beta(alpha, alpha) batch_size = x.size(0) index = torch.randperm(batch_size) mixed_x = lam * x + (1 - lam) * x[index] return mixed_x, y, y[index], lam6. 模型部署建议
6.1 ONNX导出
dummy_input = torch.randn(1, 1, 2048) torch.onnx.export( model, dummy_input, "bearing_model.onnx", input_names=["vibration_signal"], output_names=["fault_prob"], dynamic_axes={ 'vibration_signal': {0: 'batch_size'}, 'fault_prob': {0: 'batch_size'} } )6.2 边缘设备优化
使用TensorRT加速:
trtexec --onnx=bearing_model.onnx \ --saveEngine=bearing_model.trt \ --fp16 \ --workspace=2048