轴承寿命预测pytorch代码

以下是基于PyTorch实现的轴承寿命预测代码：

import torch
import torch.nn as nn
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split

# 加载数据集
data = pd.read_csv('bearing_dataset.csv')

# 数据预处理
scaler = MinMaxScaler()
data = scaler.fit_transform(data)

# 切分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(data[:, :-1], data[:, -1], test_size=0.2, random_state=42)

# 转换为张量
X_train = torch.tensor(X_train).float()
y_train = torch.tensor(y_train).float()
X_test = torch.tensor(X_test).float()
y_test = torch.tensor(y_test).float()

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(4, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 1)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(p=0.2)
        
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc3(x)
        return x

# 训练模型
net = Net()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(net.parameters(), lr=0.01)

for epoch in range(1000):
    optimizer.zero_grad()
    outputs = net(X_train)
    loss = criterion(outputs, y_train.unsqueeze(1))
    loss.backward()
    optimizer.step()
    if epoch % 100 == 0:
        print('Epoch {}, Loss: {:.4f}'.format(epoch, loss.item()))

# 测试模型
with torch.no_grad():
    outputs = net(X_test)
    loss = criterion(outputs, y_test.unsqueeze(1))
    print('Test Loss: {:.4f}'.format(loss.item()))

# 保存模型
torch.save(net.state_dict(), 'bearing_model.pt')

注意，这里使用的是简单的三层全连接神经网络，输入特征为4维，即四个传感器采集的数据。在训练过程中，使用均方误差作为损失函数，Adam优化器进行梯度下降。最终将训练好的模型保存到`bearing_model.pt`文件中。