用pytorch写一段代码：以西储大学CWRU轴承故障的振动信号数据库作为模型的训练集，故障类型标签包括：正常、内圈故障、外圈故障和滚动体故障四类。首先将一维故障信号转换为二维灰度图作为CNN模型的输...

下面是一段使用PyTorch构建CNN模型，使用CWRU轴承故障振动信号数据库进行训练的代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from torchvision.datasets import ImageFolder

# 定义CNN模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(64 * 6 * 6, 256)
        self.fc2 = nn.Linear(256, 4)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = nn.functional.relu(x)
        x = self.pool2(x)
        x = self.conv3(x)
        x = nn.functional.relu(x)
        x = self.pool3(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.fc2(x)
        return x

# 加载数据集
train_dataset = ImageFolder("path/to/training/dataset", transform=ToTensor())
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 定义训练参数
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(CNN().parameters(), lr=0.001)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = CNN()(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        if (i+1) % 10 == 0:
            print('Epoch [{}/{}], Batch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, i+1, len(train_loader), loss.item()))

在这里，我们使用了PyTorch的`nn.Module`类来定义CNN模型，然后使用`ImageFolder`类来加载CWRU轴承故障振动信号数据库，将一维故障信号转换为二维灰度图作为CNN模型的输入。我们还使用`nn.CrossEntropyLoss()`作为损失函数，使用`optim.Adam()`作为优化器，最后训练模型。