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

引言

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

准备工作

安装依赖：
确保你已经安装了以下 Python 库：
```
pip install pandas numpy torch matplotlib scikit-learn
```
下载数据：
使用 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()

使用说明

保存脚本：
- 将训练脚本代码保存为 train.py。
- 将模型定义代码保存为 model.py。
- 将测试脚本代码保存为 test.py。
运行训练：
- 在命令行中运行训练脚本：
```
python train.py
```
- 训练完成后，模型将保存为 lstm_model.pth。
运行测试和可视化：
- 在命令行中运行测试脚本：
```
python test.py
```
- 这将加载已训练的模型，并可视化预测结果。
  
  这只是一个演示，模型的预测效果还有待进一步优化。

模型调整

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

数据规模不足：
- 如果训练数据不足，模型可能无法学到市场的长期趋势。
- 改进：使用更多的历史数据，尽量包括多年的数据。可以尝试增加数据的时间跨度。
数据预处理问题：
- 数据没有正确归一化，或归一化范围过窄。
- 改进：检查 MinMaxScaler 的应用。你可以尝试不同的归一化范围，例如 (0, 1) 或 (-1, 1)，也可以使用其他标准化方法（例如 StandardScaler）。
模型复杂度不足：
- 模型的层数或隐藏单元数量可能不足以捕捉数据的复杂性。
- 改进：增加 LSTM 的隐藏层数量或隐藏单元数量。你还可以考虑添加其他类型的层，例如卷积层（CNN）或全连接层，以提高模型的表达能力。
超参数调整：
- 学习率、批大小和时间步长等超参数可能需要调整以优化模型性能。
- 改进：尝试不同的学习率（例如，0.001、0.0001 等）、不同的批大小（如 16、32、64）和时间步长（如 30、60）。
更改损失函数：
- 在某些情况下，使用不同的损失函数可能有助于模型的收敛。
- 改进：可以尝试使用其他损失函数，例如 Huber 损失函数（nn.SmoothL1Loss）或自定义损失函数，以更好地适应数据。