使用pytroch用LSTm进行回归预测代码
时间: 2024-02-11 21:09:02 浏览: 21
以下是使用PyTorch实现LSTM进行回归预测的代码示例:
```python
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
# 生成数据
data_len = 1000
time_steps = 20
x = np.arange(data_len)
y = np.sin(x * (2 * np.pi / time_steps))
# 数据预处理
x = x.reshape(-1, time_steps)
y = y.reshape(-1, time_steps)
x = torch.from_numpy(x).float()
y = torch.from_numpy(y).float()
# 定义LSTM模型
class LSTM(nn.Module):
def __init__(self, input_size, hidden_size, num_layers, output_size):
super(LSTM, self).__init__()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
self.fc = nn.Linear(hidden_size, output_size)
def forward(self, x):
h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
out, _ = self.lstm(x, (h0, c0))
out = out[:, -1, :]
out = self.fc(out)
return out
# 定义模型训练函数
def train(model, x, y, num_epochs, learning_rate):
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
losses = []
for epoch in range(num_epochs):
outputs = model(x)
loss = criterion(outputs, y)
losses.append(loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (epoch+1) % 100 == 0:
print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))
return losses
# 设置超参数
input_size = 1
hidden_size = 32
num_layers = 2
output_size = 1
num_epochs = 1000
learning_rate = 0.01
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# 创建模型实例并进行训练
model = LSTM(input_size, hidden_size, num_layers, output_size).to(device)
losses = train(model, x, y, num_epochs, learning_rate)
# 绘制损失函数曲线
plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.show()
# 使用训练好的模型进行预测并绘制结果
with torch.no_grad():
test_x = torch.from_numpy(np.arange(1000, 1200)).float().reshape(-1, time_steps).to(device)
test_y = model(test_x)
plt.plot(test_x.cpu().numpy().flatten(), test_y.cpu().numpy().flatten(), label='Predicted')
plt.plot(x.numpy().flatten(), y.numpy().flatten(), label='Ground Truth')
plt.legend()
plt.show()
```
在上述代码中,我们首先生成了一个sin函数的数据集,然后定义了一个LSTM模型,并使用该模型对数据进行训练。最后,我们使用训练好的模型对测试集进行预测,并绘制了预测结果和真实值的对比图。