pytorch 的lstm实现时间序列预测的例子
时间: 2024-01-21 17:03:19 浏览: 87
很高兴回答你的问题!以下是一个使用PyTorch LSTM实现时间序列预测的例子:
```
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
# 准备数据
data = np.sin(np.arange(1000)*(np.pi/32))
train_x, train_y = [], []
n_step = 20
for i in range(len(data)-n_step-1):
x = data[i:i+n_step]
y = data[i+n_step]
train_x.append(x)
train_y.append(y)
train_x = np.array(train_x)[..., np.newaxis]
train_y = np.array(train_y)[..., np.newaxis]
# 定义模型
class LSTM(nn.Module):
def __init__(self, input_size=1, hidden_size=10, output_size=1, num_layers=2):
super().__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=None):
out, (hn, cn) = self.lstm(x, h0)
out = self.fc(out[:, -1, :])
return out, (hn, cn)
# 训练模型
model = LSTM()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
h0 = None
batch_size = 32
n_epoch = 200
for epoch in range(n_epoch):
for i in range(0, len(train_x), batch_size):
batch_x = torch.tensor(train_x[i:i+batch_size], dtype=torch.float32)
batch_y = torch.tensor(train_y[i:i+batch_size], dtype=torch.float32)
optimizer.zero_grad()
if h0 is not None:
h0 = (torch.zeros((model.num_layers, batch_size, model.hidden_size)),
torch.zeros((model.num_layers, batch_size, model.hidden_size)))
outputs, _ = model(batch_x, h0)
loss = criterion(outputs, batch_y)
loss.backward()
optimizer.step()
if epoch % 10 == 0:
print('Epoch {}/{}, Loss: {:.6f}'.format(epoch+1, n_epoch, loss))
# 测试模型
model.eval()
test_x = train_x[200].reshape(1, -1, 1)
test_y = train_y[200]
y_pred = []
for i in range(20):
test_x = torch.tensor(test_x, dtype=torch.float32)
output, h0 = model(test_x)
test_x = output.detach().numpy()
y_pred.append(output.detach().numpy()[0][0])
test_x = np.append(test_x, train_x[200][i+1:])
test_x = test_x.reshape(1, -1, 1)
# 可视化结果
plt.plot(train_y[200:220], label='true')
plt.plot(y_pred, label='predicted')
plt.legend()
plt.show()
```
该模型用于预测一个正弦波样本序列。模型使用了一个LSTM层和一个全连接层。在训练过程中,将数据分成了多个batch,并使用MSE损失函数和Adam优化器进行优化。训练完成后,我们可以使用模型预测未来的时间序列。最后,我们将真实值和预测值绘制在一张图中进行可视化。
希望这个例子能够帮助到你!如果你有任何其他问题,请随时向我提出。
阅读全文
相关推荐


















