用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()
登录后复制