用gru代码对美国凯斯西储大学轴承故障数据集进行故障诊断
时间: 2023-09-05 08:10:48 浏览: 138
以下是使用 PyTorch 实现 GRU 模型对美国凯斯西储大学轴承故障数据集进行故障诊断的示例代码:
```python
import torch
import torch.nn as nn
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
# 读取数据
data = pd.read_csv('K001_1_48k_0_12hp_B_C3_0.csv', header=None)
data = data.values[:, 1]
data = data.reshape(-1, 1)
# 数据预处理
scaler = MinMaxScaler(feature_range=(-1, 1))
data_scaled = scaler.fit_transform(data)
# 定义超参数
input_size = 50
hidden_size = 32
num_layers = 2
output_size = 1
learning_rate = 0.001
num_epochs = 100
# 创建训练数据集
X = []
y = []
for i in range(input_size, len(data_scaled)):
X.append(data_scaled[i-input_size:i])
y.append(data_scaled[i])
X, y = np.array(X), np.array(y)
# 划分训练集和测试集
train_size = int(len(data_scaled) * 0.8)
X_train, X_test = X[:train_size], X[train_size:]
y_train, y_test = y[:train_size], y[train_size:]
# 创建数据加载器
train_data = torch.utils.data.TensorDataset(torch.from_numpy(X_train).float(), torch.from_numpy(y_train).float())
train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True)
# 定义 GRU 模型
class GRU(nn.Module):
def __init__(self, input_size, hidden_size, num_layers, output_size):
super(GRU, self).__init__()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.gru = nn.GRU(input_size, hidden_size, num_layers, batch_first=True)
self.fc = nn.Linear(hidden_size, output_size)
def forward(self, x, h0):
out, h = self.gru(x, h0)
out = self.fc(out[:, -1, :])
return out, h
model = GRU(input_size, hidden_size, num_layers, output_size)
# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# 训练模型
for epoch in range(num_epochs):
for i, (inputs, labels) in enumerate(train_loader):
# 前向传播
outputs, _ = model(inputs, None)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i+1) % 10 == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, i+1, len(train_loader), loss.item()))
# 预测测试集
with torch.no_grad():
inputs = torch.from_numpy(X_test).float()
outputs, _ = model(inputs, None)
predicted = scaler.inverse_transform(outputs.numpy())
actual = scaler.inverse_transform(y_test)
# 计算 R2 得分
from sklearn.metrics import r2_score
print('R2 score:', r2_score(actual, predicted))
```
在上述代码中,我们首先读取了凯斯西储大学轴承故障数据集,并进行了数据预处理。然后,我们定义了超参数,包括输入大小、隐藏层大小、层数、输出大小、学习率和训练轮数。接着,我们创建了训练数据集和数据加载器,并定义了 GRU 模型、损失函数和优化器。在训练循环中,我们对每个 mini-batch 执行前向传播、损失计算、反向传播和优化。在预测测试集时,我们使用训练好的模型进行前向传播,并将预测结果反归一化。最后,我们计算模型在测试集上的 R2 得分。
阅读全文