基于 pytorch + lstm 进行时间序列预测(附完整源码)

基于PyTorch LSTM（长短期记忆网络）进行时间序列预测的代码如下所示：

import torch
import torch.nn as nn
import numpy as np

# 定义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 = self.fc(out[:, -1, :])
        return out

# 设置随机种子，以便复现结果
torch.manual_seed(42)

# 超参数
input_size = 1
hidden_size = 32
num_layers = 2
output_size = 1
num_epochs = 100
learning_rate = 0.001

# 创建数据集(假设我们有一个包含100个数据点的时间序列)
data = np.sin(np.arange(0, 10*np.pi, 0.1))
data = data[:, None]

# 划分训练集和测试集
train_size = int(len(data) * 0.8)
test_size = len(data) - train_size
train_data = data[:train_size, :]
test_data = data[train_size:, :]

# 将数据集转化为PyTorch的张量
train_data = torch.Tensor(train_data).float()
test_data = torch.Tensor(test_data).float()

# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化模型
model = LSTM(input_size, hidden_size, num_layers, output_size).to(device)

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(num_epochs):
    model.train()
    
    # 前向传播和计算损失
    output = model(train_data)
    loss = criterion(output, train_data)
    
    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if (epoch+1) % 10 == 0:
        print(f"Epoch: {epoch+1}, Loss: {loss.item()}")

# 在测试集上进行预测
model.eval()
with torch.no_grad():
  # 训练数据结果
  train_result = model(train_data).detach().numpy()
  # 测试数据结果
  test_result = model(test_data).detach().numpy()

# 绘制结果图
import matplotlib.pyplot as plt

plt.figure(figsize=(12, 8))
plt.plot(train_data.numpy(), label='Original data')
plt.plot(range(train_size, len(data)), train_result, label='Train prediction')
plt.plot(range(train_size, len(data)), test_result, label='Test prediction')
plt.legend()
plt.show()

上述代码首先定义了一个LSTM模型，然后设置了超参数，创建了一个包含100个数据点的时间序列，并将其划分为训练集和测试集。接下来，将数据转换为PyTorch张量，并定义了设备（CPU或GPU）。

然后，初始化模型，并定义损失函数和优化器。在训练过程中，进行前向传播、计算损失、反向传播和优化，然后将损失打印出来。

在测试阶段，使用训练好的模型对训练集和测试集进行预测，并将结果可视化显示出来。最后，使用matplotlib库绘制了原始数据、训练预测和测试预测的图形。