### 使用 PyTorch 实现基于卷积神经网络（CNN）的故障诊断系统 以下是关于如何使用 PyTorch 构建一个基于 CNN 的旋转机械故障诊断系统的详细介绍： #### 1. 数据准备 为了训练和评估模型，需要准备好用于故障诊断的数据集。常见的数据集包括 CWRU 西储大学轴承数据集、江南大学轴承数据集等[^2]。这些数据通常是以振动信号的形式存在。 在加载数据之前，可能需要对其进行预处理，例如标准化、分割成固定长度的时间窗口以及转换为适合输入到 CNN 的张量形式。如果数据是一维时间序列，则可以通过 `unsqueeze` 方法增加维度以适配 CNN 输入的要求[^3]。 ```python import torch from torchvision import transforms # 定义数据变换操作 transform = transforms.Compose([ lambda x: torch.tensor(x, dtype=torch.float), lambda x: x.unsqueeze(0) # 增加通道维度 (batch_size, channels, length) ]) class FaultDataset(torch.utils.data.Dataset): def __init__(self, data, labels, transform=None): self.data = data self.labels = labels self.transform = transform def __len__(self): return len(self.data) def __getitem__(self, idx): sample = self.data[idx] label = self.labels[idx] if self.transform: sample = self.transform(sample) return sample, label ``` --- #### 2. 模型定义 构建一个多尺度卷积神经网络（MS-CNN），该网络能够提取不同尺度下的特征并结合注意力机制提升抗噪性能。下面展示了一个简单的 MS-CNN 结构示例： ```python import torch.nn as nn import torch.nn.functional as F class MSCNN(nn.Module): def __init__(self, num_classes=10): super(MSCNN, self).__init__() # 多尺度卷积层 self.conv1_1 = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=9, padding='same') self.conv1_2 = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=17, padding='same') self.conv1_3 = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=33, padding='same') # 后续共享层 self.pool = nn.MaxPool1d(kernel_size=2, stride=2) self.fc1 = nn.Linear(48 * 125, 128) # 根据输入尺寸调整全连接层参数 self.fc2 = nn.Linear(128, num_classes) def forward(self, x): # 提取多尺度特征 x1 = F.relu(self.conv1_1(x)) x2 = F.relu(self.conv1_2(x)) x3 = F.relu(self.conv1_3(x)) # 特征融合 x = torch.cat((x1, x2, x3), dim=1) x = self.pool(x) x = x.view(-1, 48 * 125) # 展平 # 全连接层 x = F.relu(self.fc1(x)) x = self.fc2(x) return x ``` --- #### 3. 训练过程 设置优化器、损失函数，并编写训练循环逻辑。这里采用交叉熵损失函数和 Adam 优化算法进行训练。 ```python model = MSCNN(num_classes=10).cuda() # 将模型移动至 GPU 上运行 criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) def train_model(dataloader, model, criterion, optimizer, epochs=10): for epoch in range(epochs): running_loss = 0.0 correct_predictions = 0 total_samples = 0 for inputs, labels in dataloader: inputs, labels = inputs.cuda(), labels.cuda() outputs = model(inputs) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() _, predicted = torch.max(outputs, 1) total_samples += labels.size(0) correct_predictions += (predicted == labels).sum().item() running_loss += loss.item() accuracy = correct_predictions / total_samples print(f'Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(dataloader):.4f}, Accuracy: {accuracy*100:.2f}%') train_loader = torch.utils.data.DataLoader(FaultDataset(data, labels, transform=transform), batch_size=32, shuffle=True) train_model(train_loader, model, criterion, optimizer, epochs=20) ``` --- #### 4. 测试与可视化 完成训练后，可以在测试集上评估模型表现，并绘制混淆矩阵或 t-SNE 图像进一步分析结果。 ```python from sklearn.metrics import confusion_matrix import seaborn as sns import matplotlib.pyplot as plt def evaluate_model(test_loader, model): y_true, y_pred = [], [] with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.cuda(), labels.cuda() outputs = model(inputs) _, preds = torch.max(outputs, 1) y_true.extend(labels.cpu().numpy()) y_pred.extend(preds.cpu().numpy()) cm = confusion_matrix(y_true, y_pred) sns.heatmap(cm, annot=True, fmt="d", cmap="Blues") plt.title('Confusion Matrix') plt.show() test_loader = torch.utils.data.DataLoader(FaultDataset(test_data, test_labels, transform=transform), batch_size=32, shuffle=False) evaluate_model(test_loader, model) ``` --- ### 总结 上述代码展示了如何使用 PyTorch 来实现一个基于 CNN 的旋转机械故障诊断系统。通过引入多尺度卷积核和注意力机制，可以有效增强模型对抗噪声的能力，从而提高诊断精度。