PyTorch LSTM 单步、多步时间预测

多维输入、多维输出；单步预测、多步滚动预测

import torch
import torch.nn as nn
import torch.optim as optimclass LSTMModel(nn.Module):def __init__(self, input_dim, hidden_dim, num_layers, output_dim):super(LSTMModel, self).__init__()self.hidden_dim = hidden_dimself.num_layers = num_layersself.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)self.fc = nn.Linear(hidden_dim, output_dim)def forward(self, x):h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)out, _ = self.lstm(x, (h0, c0))out = self.fc(out[:, -1, :])return out# 超参数
input_dim = 400
hidden_dim = 64
num_layers = 2
output_dim = 1
num_epochs = 100
learning_rate = 0.001
batch_size = 32# 初始化模型、损失函数和优化器
model = LSTMModel(input_dim, hidden_dim, num_layers, output_dim)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)# 示例训练代码（假设已经定义了train_loader）
for epoch in range(num_epochs):for i, (inputs, labels) in enumerate(train_loader):model.train()outputs = model(inputs)loss = criterion(outputs, labels)optimizer.zero_grad()loss.backward()optimizer.step()if (i+1) % 100 == 0:print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')# 保存模型
torch.save(model.state_dict(), 'lstm_model_single_step.pth')# 多步预测函数
def multi_step_predict(model, input_seq, future_steps):model.eval()  # 切换到评估模式predictions = []input_seq = input_seq.unsqueeze(0)  # 增加batch维度，shape变为 (1, seq_len, input_dim)for _ in range(future_steps):with torch.no_grad():  # 禁用梯度计算pred = model(input_seq)  # 预测下一个时间步predictions.append(pred.item())  # 存储预测值# 更新输入序列，将预测值添加到末尾，并移除最早的一个时间步input_seq = torch.cat((input_seq[:, 1:, :], pred.unsqueeze(0).unsqueeze(2)), dim=1)return predictions# 示例调用
initial_input_seq = torch.randn(1, 155, 400)  # 假设这是的初始输入
future_steps = 10
predictions = multi_step_predict(model, initial_input_seq, future_steps)
print(predictions)