用PythonPyTorch实现轴承振动数据的智能故障诊断从零搭建CNN模型轴承作为旋转机械的核心部件其健康状态直接影响设备寿命。传统振动分析依赖工程师手动提取时频域特征既耗时又难以捕捉复杂故障模式。我们以凯斯西储大学轴承数据集为例用PyTorch构建端到端的智能诊断系统实现振动信号到故障类别的自动映射。1. 环境准备与数据加载工欲善其事必先利其器。建议使用Python 3.8环境通过conda创建虚拟环境避免依赖冲突conda create -n fault_diagnosis python3.8 conda activate fault_diagnosis pip install torch torchvision pandas scikit-learn matplotlib凯斯西储大学轴承数据集包含正常状态和多种故障类型内圈、外圈、滚动体故障采样频率12kHz。我们使用pandas加载预处理后的CSV版本import pandas as pd from sklearn.model_selection import train_test_split data pd.read_csv(bearing_vibration.csv) X data.iloc[:, :-1].values # 振动信号序列 y data.iloc[:, -1].values # 故障标签 # 划分训练集和测试集 X_train, X_test, y_train, y_test train_test_split( X, y, test_size0.2, stratifyy, random_state42 )提示工业现场数据往往存在类别不平衡问题可采用过采样或加权损失函数处理2. 数据预处理与增强策略原始振动信号需要转化为适合CNN输入的格式。我们采用滑动窗口技术生成样本片段import numpy as np def create_sequences(data, labels, window_size1024, step512): sequences [] sequence_labels [] for i in range(0, len(data) - window_size, step): seq data[i:iwindow_size] sequences.append(seq) sequence_labels.append(labels[i window_size // 2]) return np.array(sequences), np.array(sequence_labels) X_train_seq, y_train_seq create_sequences(X_train, y_train) X_test_seq, y_test_seq create_sequences(X_test, y_test)为提升模型泛化能力引入三种数据增强技术随机缩放振幅±10%波动随机偏移添加±0.1倍标准差噪声随机裁剪从1024点中随机选取896点class VibrationAugment: def __call__(self, sample): if np.random.rand() 0.5: sample * 1 np.random.uniform(-0.1, 0.1) if np.random.rand() 0.5: sample np.random.normal(0, 0.1*sample.std()) if np.random.rand() 0.5: start np.random.randint(0, 128) sample sample[start:start896] sample np.pad(sample, (0, 1024-896)) return sample3. 构建1D-CNN诊断模型针对振动信号特性设计包含并行卷积支路的网络结构import torch import torch.nn as nn import torch.nn.functional as F class FaultDiagnosisCNN(nn.Module): def __init__(self, num_classes): super().__init__() # 高频特征支路 self.branch1 nn.Sequential( nn.Conv1d(1, 16, kernel_size64, stride8, padding28), nn.BatchNorm1d(16), nn.ReLU(), nn.MaxPool1d(2) ) # 低频特征支路 self.branch2 nn.Sequential( nn.Conv1d(1, 16, kernel_size256, stride32, padding112), nn.BatchNorm1d(16), nn.ReLU(), nn.MaxPool1d(2) ) # 特征融合 self.fc nn.Sequential( nn.Linear(16*2*31, 64), nn.ReLU(), nn.Dropout(0.5), nn.Linear(64, num_classes) ) def forward(self, x): x x.unsqueeze(1) # [batch, 1, seq_len] h1 self.branch1(x) h2 self.branch2(x) h torch.cat([h1, h2], dim1) return self.fc(h.view(h.size(0), -1))模型设计关键点双支路结构分别捕捉不同时间尺度的故障特征宽卷积核适应振动信号的周期性特点BatchNorm加速训练并提升稳定性Dropout防止过拟合4. 模型训练与性能优化采用带热启动的余弦退火学习率策略配合标签平滑技术from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts model FaultDiagnosisCNN(num_classeslen(np.unique(y))) criterion nn.CrossEntropyLoss(label_smoothing0.1) optimizer torch.optim.AdamW(model.parameters(), lr1e-3) scheduler CosineAnnealingWarmRestarts(optimizer, T_010, T_mult2) for epoch in range(100): model.train() for batch_x, batch_y in train_loader: optimizer.zero_grad() outputs model(batch_x) loss criterion(outputs, batch_y) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step()评估指标除准确率外还应关注混淆矩阵识别易混淆故障类型F1-score处理类别不平衡ROC-AUC综合评估分类性能from sklearn.metrics import classification_report model.eval() with torch.no_grad(): y_pred model(X_test_tensor).argmax(dim1) print(classification_report(y_test, y_pred.cpu()))5. 模型解释与工业部署使用Grad-CAM可视化关键故障特征区域class GradCAM: def __init__(self, model, target_layer): self.model model self.target_layer target_layer self.gradients None def save_gradient(self, grad): self.gradients grad def __call__(self, x): feature_maps None def hook_fn(module, input, output): nonlocal feature_maps feature_maps output output.register_hook(self.save_gradient) handle self.target_layer.register_forward_hook(hook_fn) output self.model(x) handle.remove() one_hot torch.zeros_like(output) one_hot[0, output.argmax()] 1 self.model.zero_grad() output.backward(gradientone_hot, retain_graphTrue) weights self.gradients.mean(dim2, keepdimTrue) cam (weights * feature_maps).sum(dim1, keepdimTrue) cam F.relu(cam) cam F.interpolate(cam, sizex.shape[2], modelinear) return cam.squeeze().cpu().numpy()工业部署建议ONNX导出实现跨平台部署TensorRT优化提升推理速度边缘设备适配使用LibTorch在嵌入式系统运行dummy_input torch.randn(1, 1, 1024) torch.onnx.export(model, dummy_input, fault_diagnosis.onnx, input_names[vibration], output_names[fault_class])6. 持续改进方向在实际项目中我们发现以下优化手段效果显著混合精度训练减少显存占用知识蒸馏压缩模型尺寸多传感器融合结合温度、电流信号半监督学习利用未标注数据from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): outputs model(batch_x) loss criterion(outputs, batch_y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()
别再手动提特征了！用Python+PyTorch搭建你的第一个智能故障诊断模型（以轴承振动数据为例）
发布时间：2026/6/1 18:46:09
用PythonPyTorch实现轴承振动数据的智能故障诊断从零搭建CNN模型轴承作为旋转机械的核心部件其健康状态直接影响设备寿命。传统振动分析依赖工程师手动提取时频域特征既耗时又难以捕捉复杂故障模式。我们以凯斯西储大学轴承数据集为例用PyTorch构建端到端的智能诊断系统实现振动信号到故障类别的自动映射。1. 环境准备与数据加载工欲善其事必先利其器。建议使用Python 3.8环境通过conda创建虚拟环境避免依赖冲突conda create -n fault_diagnosis python3.8 conda activate fault_diagnosis pip install torch torchvision pandas scikit-learn matplotlib凯斯西储大学轴承数据集包含正常状态和多种故障类型内圈、外圈、滚动体故障采样频率12kHz。我们使用pandas加载预处理后的CSV版本import pandas as pd from sklearn.model_selection import train_test_split data pd.read_csv(bearing_vibration.csv) X data.iloc[:, :-1].values # 振动信号序列 y data.iloc[:, -1].values # 故障标签 # 划分训练集和测试集 X_train, X_test, y_train, y_test train_test_split( X, y, test_size0.2, stratifyy, random_state42 )提示工业现场数据往往存在类别不平衡问题可采用过采样或加权损失函数处理2. 数据预处理与增强策略原始振动信号需要转化为适合CNN输入的格式。我们采用滑动窗口技术生成样本片段import numpy as np def create_sequences(data, labels, window_size1024, step512): sequences [] sequence_labels [] for i in range(0, len(data) - window_size, step): seq data[i:iwindow_size] sequences.append(seq) sequence_labels.append(labels[i window_size // 2]) return np.array(sequences), np.array(sequence_labels) X_train_seq, y_train_seq create_sequences(X_train, y_train) X_test_seq, y_test_seq create_sequences(X_test, y_test)为提升模型泛化能力引入三种数据增强技术随机缩放振幅±10%波动随机偏移添加±0.1倍标准差噪声随机裁剪从1024点中随机选取896点class VibrationAugment: def __call__(self, sample): if np.random.rand() 0.5: sample * 1 np.random.uniform(-0.1, 0.1) if np.random.rand() 0.5: sample np.random.normal(0, 0.1*sample.std()) if np.random.rand() 0.5: start np.random.randint(0, 128) sample sample[start:start896] sample np.pad(sample, (0, 1024-896)) return sample3. 构建1D-CNN诊断模型针对振动信号特性设计包含并行卷积支路的网络结构import torch import torch.nn as nn import torch.nn.functional as F class FaultDiagnosisCNN(nn.Module): def __init__(self, num_classes): super().__init__() # 高频特征支路 self.branch1 nn.Sequential( nn.Conv1d(1, 16, kernel_size64, stride8, padding28), nn.BatchNorm1d(16), nn.ReLU(), nn.MaxPool1d(2) ) # 低频特征支路 self.branch2 nn.Sequential( nn.Conv1d(1, 16, kernel_size256, stride32, padding112), nn.BatchNorm1d(16), nn.ReLU(), nn.MaxPool1d(2) ) # 特征融合 self.fc nn.Sequential( nn.Linear(16*2*31, 64), nn.ReLU(), nn.Dropout(0.5), nn.Linear(64, num_classes) ) def forward(self, x): x x.unsqueeze(1) # [batch, 1, seq_len] h1 self.branch1(x) h2 self.branch2(x) h torch.cat([h1, h2], dim1) return self.fc(h.view(h.size(0), -1))模型设计关键点双支路结构分别捕捉不同时间尺度的故障特征宽卷积核适应振动信号的周期性特点BatchNorm加速训练并提升稳定性Dropout防止过拟合4. 模型训练与性能优化采用带热启动的余弦退火学习率策略配合标签平滑技术from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts model FaultDiagnosisCNN(num_classeslen(np.unique(y))) criterion nn.CrossEntropyLoss(label_smoothing0.1) optimizer torch.optim.AdamW(model.parameters(), lr1e-3) scheduler CosineAnnealingWarmRestarts(optimizer, T_010, T_mult2) for epoch in range(100): model.train() for batch_x, batch_y in train_loader: optimizer.zero_grad() outputs model(batch_x) loss criterion(outputs, batch_y) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step()评估指标除准确率外还应关注混淆矩阵识别易混淆故障类型F1-score处理类别不平衡ROC-AUC综合评估分类性能from sklearn.metrics import classification_report model.eval() with torch.no_grad(): y_pred model(X_test_tensor).argmax(dim1) print(classification_report(y_test, y_pred.cpu()))5. 模型解释与工业部署使用Grad-CAM可视化关键故障特征区域class GradCAM: def __init__(self, model, target_layer): self.model model self.target_layer target_layer self.gradients None def save_gradient(self, grad): self.gradients grad def __call__(self, x): feature_maps None def hook_fn(module, input, output): nonlocal feature_maps feature_maps output output.register_hook(self.save_gradient) handle self.target_layer.register_forward_hook(hook_fn) output self.model(x) handle.remove() one_hot torch.zeros_like(output) one_hot[0, output.argmax()] 1 self.model.zero_grad() output.backward(gradientone_hot, retain_graphTrue) weights self.gradients.mean(dim2, keepdimTrue) cam (weights * feature_maps).sum(dim1, keepdimTrue) cam F.relu(cam) cam F.interpolate(cam, sizex.shape[2], modelinear) return cam.squeeze().cpu().numpy()工业部署建议ONNX导出实现跨平台部署TensorRT优化提升推理速度边缘设备适配使用LibTorch在嵌入式系统运行dummy_input torch.randn(1, 1, 1024) torch.onnx.export(model, dummy_input, fault_diagnosis.onnx, input_names[vibration], output_names[fault_class])6. 持续改进方向在实际项目中我们发现以下优化手段效果显著混合精度训练减少显存占用知识蒸馏压缩模型尺寸多传感器融合结合温度、电流信号半监督学习利用未标注数据from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): outputs model(batch_x) loss criterion(outputs, batch_y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()
