使用 PyTorch 构建 LSTM 股票价格预测模型

引言

在金融领域，股票价格预测是一个重要且具有挑战性的任务。随着深度学习的发展，长短期记忆网络（LSTM）因其在处理时间序列数据方面的出色表现而受到关注。本篇博客将指导你如何使用PyTorch构建一个LSTM模型来预测股票价格，我们将逐步介绍数据预处理、模型训练和结果可视化的完整流程。

准备工作

1. 安装依赖：
   确保你已经安装了以下 Python 库：

   pip install pandas numpy torch matplotlib scikit-learn
   

2. 下载数据：
   使用 yfinance 库下载你感兴趣的股票的历史数据，并保存为 CSV 文件。我们这里使用 Apple（AAPL）过去五年的数据，文件命名为 AAPL_5y_data.csv。以下是一个下载数据的代码示例：

   import yfinance as yf# 下载Apple股票过去5年的数据
   data = yf.download('AAPL', start='2019-01-01', end='2024-01-01')
   data.to_csv('AAPL_5y_data.csv')
   

1. 训练模型（train.py）

在这个脚本中，我们将读取 CSV 文件，归一化数据，并使用 LSTM 模型进行训练。

import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from sklearn.preprocessing import MinMaxScaler
from model import LSTM  # 导入LSTM类# 设置随机种子
torch.manual_seed(42)# 读取CSV文件
file_path = 'AAPL_5y_data.csv'  # 替换为你的CSV文件路径
data = pd.read_csv(file_path)# 确保日期列是 datetime 类型
data['Date'] = pd.to_datetime(data['Date'])
data.set_index('Date', inplace=True)# 选择多特征：'Close', 'Open', 'High', 'Low', 'Volume'
features = data[['Close', 'Open', 'High', 'Low', 'Volume']].values# 数据归一化
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(features)# 准备训练和测试数据
train_size = int(len(scaled_data) * 0.8)
train_data = scaled_data[:train_size]
test_data = scaled_data[train_size:]def create_dataset(data, time_step=1):X, y = [], []for i in range(len(data) - time_step - 1):a = data[i:(i + time_step)]X.append(a)y.append(data[i + time_step, 0])  # 预测收盘价return np.array(X), np.array(y)# 创建数据集
time_step = 50  # 时间步长
X_train, y_train = create_dataset(train_data, time_step)# 转换为PyTorch张量
X_train = torch.from_numpy(X_train).float()
y_train = torch.from_numpy(y_train).float().view(-1, 1)# 初始化模型、损失函数和优化器
model = LSTM()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)# 训练模型
num_epochs = 300
for epoch in range(num_epochs):model.train()optimizer.zero_grad()outputs = model(X_train)loss = criterion(outputs, y_train)loss.backward()optimizer.step()if (epoch + 1) % 10 == 0:print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')# 保存模型
torch.save(model.state_dict(), 'lstm_model.pth')
print("模型已保存为 'lstm_model.pth'")

2. 模型定义（model.py）

在这个文件中定义 LSTM 模型结构。

import torch
import torch.nn as nnclass LSTM(nn.Module):def __init__(self):super(LSTM, self).__init__()self.lstm = nn.LSTM(input_size=5, hidden_size=100, num_layers=2, batch_first=True)self.fc = nn.Linear(100, 1)def forward(self, x):out, _ = self.lstm(x)out = self.fc(out[:, -1, :])  # 取最后时间步的输出return out

3. 测试模型和可视化（test.py）

在这个脚本中，我们将加载训练好的模型，并使用测试数据进行预测和可视化。

import pandas as pd
import numpy as np
import torch
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from model import LSTM  # 导入LSTM类# 设置字体为SimHei，用于显示中文
plt.rcParams['font.family'] = 'SimHei'# 读取CSV文件
file_path = 'AAPL_5y_data.csv'  # 替换为你的CSV文件路径
data = pd.read_csv(file_path)# 确保日期列是 datetime 类型
data['Date'] = pd.to_datetime(data['Date'])
data.set_index('Date', inplace=True)# 选择多特征：'Close', 'Open', 'High', 'Low', 'Volume'
features = data[['Close', 'Open', 'High', 'Low', 'Volume']].values# 数据归一化
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(features)# 准备训练和测试数据
train_size = int(len(scaled_data) * 0.8)
train_data = scaled_data[:train_size]
test_data = scaled_data[train_size:]def create_dataset(data, time_step=1):X, y = [], []for i in range(len(data) - time_step - 1):a = data[i:(i + time_step)]X.append(a)y.append(data[i + time_step, 0])  # 预测收盘价return np.array(X), np.array(y)# 创建测试数据集
time_step = 50  # 时间步长
X_test, y_test = create_dataset(test_data, time_step)# 转换为PyTorch张量
X_test = torch.from_numpy(X_test).float()
y_test = torch.from_numpy(y_test).float().view(-1, 1)# 加载模型
model = LSTM()
model.load_state_dict(torch.load('lstm_model.pth'))
model.eval()# 测试模型
with torch.no_grad():test_outputs = model(X_test)# test_outputs 是预测的收盘价，将其重新归一化为原始价格test_outputs = scaler.inverse_transform(np.concatenate((test_outputs.numpy(), np.zeros((test_outputs.shape[0], 4))), axis=1))[:, 0]  # 反归一化收盘价y_test_inverse = scaler.inverse_transform(np.concatenate((y_test.numpy(), np.zeros((y_test.shape[0], 4))), axis=1))[:, 0]# 可视化结果
plt.figure(figsize=(14, 7))
plt.plot(data.index[-len(y_test):], y_test_inverse, label='真实价格', color='blue')
plt.plot(data.index[-len(test_outputs):], test_outputs, label='预测价格', color='red')
plt.title('股票价格预测')
plt.xlabel('日期')
plt.ylabel('价格')
plt.legend()
plt.show()

使用说明

1. 保存脚本：

   • 将训练脚本代码保存为 train.py。
   • 将模型定义代码保存为 model.py。
   • 将测试脚本代码保存为 test.py。

2. 运行训练：

   • 在命令行中运行训练脚本：

     python train.py
     

   • 训练完成后，模型将保存为 lstm_model.pth。

3. 运行测试和可视化：

   • 在命令行中运行测试脚本：

     python test.py
     

   • 这将加载已训练的模型，并可视化预测结果。
     
     这只是一个演示，模型的预测效果还有待进一步优化。

模型调整

如果预测的价格和真实价格差距较大，可能是由于以下几个原因：

1. 数据规模不足：

   • 如果训练数据不足，模型可能无法学到市场的长期趋势。
   • 改进：使用更多的历史数据，尽量包括多年的数据。可以尝试增加数据的时间跨度。

2. 数据预处理问题：

   • 数据没有正确归一化，或归一化范围过窄。
   • 改进：检查 MinMaxScaler 的应用。你可以尝试不同的归一化范围，例如 (0, 1) 或 (-1, 1)，也可以使用其他标准化方法（例如 StandardScaler）。

3. 模型复杂度不足：

   • 模型的层数或隐藏单元数量可能不足以捕捉数据的复杂性。
   • 改进：增加 LSTM 的隐藏层数量或隐藏单元数量。你还可以考虑添加其他类型的层，例如卷积层（CNN）或全连接层，以提高模型的表达能力。

4. 超参数调整：

   • 学习率、批大小和时间步长等超参数可能需要调整以优化模型性能。
   • 改进：尝试不同的学习率（例如，0.001、0.0001 等）、不同的批大小（如 16、32、64）和时间步长（如 30、60）。

5. 更改损失函数：

   • 在某些情况下，使用不同的损失函数可能有助于模型的收敛。
   • 改进：可以尝试使用其他损失函数，例如 Huber 损失函数（nn.SmoothL1Loss）或自定义损失函数，以更好地适应数据。

结论

通过使用 PyTorch 构建 LSTM 模型，我们成功地实现了股票价格的预测。在这个过程中，我们学习了如何处理时间序列数据，构建和训练深度学习模型，以及如何评估和可视化预测结果。尽管模型的性能可能需要进一步的优化和调整，但这个示例为未来的工作奠定了基础。

希望这篇博客能够帮助你在股票价格预测方面取得更好的成果。欢迎分享你的成果和经验，或者提出你的问题！