### 使用PyTorch构建和训练CNN模型以处理轴承数据集 #### 模型构建 为了使用PyTorch构建CNN模型，首先需要定义模型的架构。通常情况下，CNN模型由多个卷积层、池化层和全连接层组成。对于轴承数据集（如CWRU），可以通过以下方式设计模型： ```python import torch import torch.nn as nn import torch.optim as optim class CNNModel(nn.Module): def __init__(self, num_classes=10): # 假设分类类别数为10 super(CNNModel, self).__init__() # 卷积层 self.conv1 = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1) self.relu1 = nn.ReLU() self.pool1 = nn.MaxPool1d(kernel_size=2, stride=2) self.conv2 = nn.Conv1d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1) self.relu2 = nn.ReLU() self.pool2 = nn.MaxPool1d(kernel_size=2, stride=2) # 全连接层 self.fc1 = nn.Linear(32 * (input_length // 4), 128) # input_length需根据实际数据长度调整 self.relu_fc = nn.ReLU() self.dropout = nn.Dropout(p=0.5) self.fc2 = nn.Linear(128, num_classes) def forward(self, x): x = self.pool1(self.relu1(self.conv1(x))) x = self.pool2(self.relu2(self.conv2(x))) x = x.view(x.size(0), -1) # 展平操作 x = self.relu_fc(self.fc1(x)) x = self.dropout(x) x = self.fc2(x) return x ``` 此代码片段展示了如何创建一个简单的CNN模型[^1]。 #### 数据预处理 在训练之前，需要对轴承数据集进行预处理。这一步骤可能包括去噪、归一化以及划分训练集、验证集和测试集。以下是数据加载和预处理的一个示例： ```python from torchvision import datasets, transforms from torch.utils.data import DataLoader transform = transforms.Compose([ transforms.ToTensor(), # 转换为张量 transforms.Normalize((mean,), (std,)) # 归一化，均值和标准差需预先计算 ]) train_dataset = CustomDataset(train_data_path, transform=transform) val_dataset = CustomDataset(val_data_path, transform=transform) test_dataset = CustomDataset(test_data_path, transform=transform) train_loader = DataLoader(dataset=train_dataset, batch_size=32, shuffle=True) val_loader = DataLoader(dataset=val_dataset, batch_size=32, shuffle=False) test_loader = DataLoader(dataset=test_dataset, batch_size=32, shuffle=False) ``` 这里假设`CustomDataset`是一个自定义的数据集类，继承自`torch.utils.data.Dataset`[^3]。 #### 训练过程 完成模型定义和数据准备后，可以开始训练模型。下面是一段完整的训练流程代码： ```python device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = CNNModel(num_classes).to(device) criterion = nn.CrossEntropyLoss() # 定义损失函数 optimizer = optim.Adam(model.parameters(), lr=0.001) # 定义优化器 num_epochs = 50 for epoch in range(num_epochs): model.train() running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader)}") print("Training complete.") ``` 这段代码实现了模型的训练循环，并打印每轮迭代的平均损失值[^4]。 #### 测试与评估 最后，在测试集上评估模型性能。可以使用混淆矩阵或其他指标衡量模型效果： ```python correct = 0 total = 0 model.eval() with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() accuracy = 100 * correct / total print(f'Test Accuracy of the model on the test set: {accuracy}%') ``` 以上代码完成了模型的测试阶段，并输出最终的准确率[^4]。 ---