1.读取数据集
import numpy as np
import torch
from torch import nn
import matplotlib.pyplot as plt
import pandas as pd
from torch.autograd import Variable
all_date = pd.read_csv(u"D:/故障诊断数据集/西交轴承数据/XJTU-SY_Bearing_Datasets/Data/XJTU-SY_Bearing_Datasets/35Hz12kN/Bearing1_1/all.csv")
all_date['RUL'] = np.arange(0, 1,1/len(all_date))
all_dateNd = all_date.values
2.搭建lstm网络
class LstmRNN(nn.Module):
"""
Parameters:
- input_size: feature size
- hidden_size: number of hidden units
- output_size: number of output
- num_layers: layers of LSTM to stack
"""
def __init__(self, input_size, hidden_size=1, output_size=4, num_layers=1):
super().__init__()
self.lstm = nn.LSTM(input_size, hidden_size, num_layers) # utilize the LSTM model in torch.nn
self.forwardCalculation = nn.Linear(hidden_size, output_size)
def forward(self, _x):
x, _ = self.lstm(_x) # _x is input, size (seq_len, batch, input_size)
s, b, h = x.shape # x is output, size (seq_len, batch, hidden_size)
x = x.view(s * b, h)
x = self.forwardCalculation(x)
x = x.view(s, b, -1)
return x
3.数据处理和网络训练文章来源:https://www.toymoban.com/news/detail-645300.html
if __name__ == '__main__':
# create database
data_len = len(all_dateNd)
data_x = all_dateNd[:,0:2] #二维输入
data_y = all_dateNd[:,2] #输出
dataset = np.zeros((data_len, 3)) #用一个dataset表示输入输出
dataset[:, 0] = data_x[:,0]
dataset[:, 1] = data_x[:,1]
dataset[:, 2] = data_y
dataset = dataset.astype('float32')
# choose dataset for training and testing
train_data_ratio = 0.1 # 数据量太大,仅使用10%做训练
train_data_len = int(data_len * train_data_ratio)
train_x = dataset[:train_data_len, 0:2]
train_y = dataset[:train_data_len, 2]
INPUT_FEATURES_NUM = 2 #输入是2维:水平和垂直振动信号
OUTPUT_FEATURES_NUM = 1 #输出是一维:寿命
# test_x = dataset[train_data_len:, 0:2]
# test_y = dataset[train_data_len:, 2]
# ----------------- train -------------------
train_x_tensor = train_x.reshape(-1, 14, INPUT_FEATURES_NUM) # 每14个做一个batch
train_y_tensor = train_y.reshape(-1, 14, OUTPUT_FEATURES_NUM) # set batch size to 14
# transfer data to pytorch tensor
train_x_tensor = torch.from_numpy(train_x_tensor) #数组转tensor
train_y_tensor = torch.from_numpy(train_y_tensor)
train_x_tensor_GPU = train_x_tensor.cuda() #用GPU训练
train_y_tensor_GPU = train_y_tensor.cuda()
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") #用GPU训练网络
lstm_model = LstmRNN(INPUT_FEATURES_NUM, 16, output_size=OUTPUT_FEATURES_NUM, num_layers=1).to(DEVICE) # 16 hidden units
print('LSTM model:', lstm_model)
print('model.parameters:', lstm_model.parameters)
loss_function = nn.MSELoss()
optimizer = torch.optim.Adam(lstm_model.parameters(), lr=0.01)
max_epochs = 100000
for epoch in range(max_epochs):
output = lstm_model(train_x_tensor_GPU)
loss = loss_function(output, train_y_tensor_GPU)
loss.backward()
optimizer.step()
optimizer.zero_grad()
if loss.item() < 1e-4:
print('Epoch [{}/{}], Loss: {:.5f}'.format(epoch + 1, max_epochs, loss.item()))
print("The loss value is reached")
output = output.view(-1, OUTPUT_FEATURES_NUM).data.cpu().numpy()
plt.plot(output, 'b')
break
elif (epoch + 1) % 100 == 0:
print('Epoch: [{}/{}], Loss:{:.5f}'.format(epoch + 1, max_epochs, loss.item()))
# prediction on training dataset
predictive_y_for_training = lstm_model(train_x_tensor_GPU)
predictive_y_for_training = predictive_y_for_training.view(-1, OUTPUT_FEATURES_NUM).data.cpu().numpy()
4.绘图文章来源地址https://www.toymoban.com/news/detail-645300.html
predictive_y_for_training = lstm_model(train_x_tensor_GPU)
predictive_y_for_training = predictive_y_for_training.view(-1, OUTPUT_FEATURES_NUM).data.cpu().numpy() #tensor转numpy数组
plt.figure()
plt.plot(train_y, 'b')
plt.plot(output, 'y--')
plt.plot(predictive_y_for_training, 'y--')
plt.show()
到了这里,关于lstm轴承寿命预测的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!