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

以下是用PyTorch实现将一维故障信号转换为二维灰度图并训练CNN模型的代码：

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split

# 加载数据集
data = np.load('CWRU_data.npy')
labels = np.load('CWRU_labels.npy')

# 将一维信号转换为二维灰度图
def convert_to_image(data):
    image = np.zeros((len(data), 100, 100))
    for i in range(len(data)):
        image[i] = np.reshape(data[i], (100, 100))
    return image

# 数据预处理
images = convert_to_image(data)
images = images.reshape(images.shape[0], 1, images.shape[1], images.shape[2])
labels = labels - 1  # 将标签从1,2,3,4转换为0,1,2,3
x_train, x_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, random_state=42)

# 定义CNN模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 5)
        self.fc1 = nn.Linear(64 * 21 * 21, 512)
        self.fc2 = nn.Linear(512, 4)

    def forward(self, x):
        x = self.pool(nn.functional.relu(self.conv1(x)))
        x = self.pool(nn.functional.relu(self.conv2(x)))
        x = x.view(-1, 64 * 21 * 21)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 训练CNN模型
model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

epochs = 20
train_loss = []
test_loss = []
train_acc = []
test_acc = []

for epoch in range(epochs):
    running_loss = 0.0
    running_acc = 0.0
    for i, data in enumerate(zip(x_train, y_train)):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs.float())
        loss = criterion(outputs, labels.long())
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        running_acc += (predicted == labels).sum().item() / len(labels)

    train_loss.append(running_loss / len(x_train))
    train_acc.append(running_acc / len(x_train))

    test_running_loss = 0.0
    test_running_acc = 0.0
    with torch.no_grad():
        for i, data in enumerate(zip(x_test, y_test)):
            inputs, labels = data
            outputs = model(inputs.float())
            loss = criterion(outputs, labels.long())
            test_running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            test_running_acc += (predicted == labels).sum().item() / len(labels)
    test_loss.append(test_running_loss / len(x_test))
    test_acc.append(test_running_acc / len(x_test))

    print("Epoch: {}/{} | Train Loss: {:.4f} | Train Acc: {:.4f} | Test Loss: {:.4f} | Test Acc: {:.4f}".format(
        epoch + 1, epochs, train_loss[-1], train_acc[-1], test_loss[-1], test_acc[-1]))

# 可视化训练过程
fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(15, 5))
ax[0].plot(train_loss, label='train loss')
ax[0].plot(test_loss, label='test loss')
ax[0].set_xlabel('Epoch')
ax[0].set_ylabel('Loss')
ax[0].legend()

ax[1].plot(train_acc, label='train accuracy')
ax[1].plot(test_acc, label='test accuracy')
ax[1].set_xlabel('Epoch')
ax[1].set_ylabel('Accuracy')
ax[1].legend()

plt.show()