LSTM在股价数据集上的预测实战
- 使用完整的JPX赛题数据，并向大家提供完整的lstm流程。
导包
import numpy as np #数据处理
import pandas as pd #数据处理
import matplotlib as mlp
import matplotlib.pyplot as plt #绘图
from sklearn.preprocessing import MinMaxScaler #·数据预处理
from sklearn.metrics import mean_squared_error
import torch
import torch.nn as nn #导入pytorch中的基本类
from torch.autograd import Variable
from torch.utils.data import DataLoader, TensorDataset
import torch.optim as optim
import torch.utils.data as data
# typing 模块提供了一些类型，辅助函数中的参数类型定义
from typing import Union,List,Tuple,Iterable
from sklearn.preprocessing import LabelEncoder,MinMaxScaler
from decimal import ROUND_HALF_UP, Decimal
一、数据加载与处理
# 一、数据加载与处理
# 1、查看数据集信息
stock= pd.read_csv('stock_prices.csv') # (2332531,12)
stock_list = pd.read_csv('stock_list.csv') # (4417,16)
stock["SecuritiesCode"].unique().__len__() #2000支股票
# 2、为了效率我们抽取其中的10支股票
selected_codes = stock['SecuritiesCode'].drop_duplicates().sample(n=10)
stock = stock[stock['SecuritiesCode'].isin(selected_codes)] # (9833,12)
stock["SecuritiesCode"].unique().__len__() #只有10支股票了
stock.isnull().sum() #查看缺失值
# 3、预处理数据集
#将Target名字修改为Sharpe Ratio
stock.rename(columns={'Target': 'Sharpe Ratio'}, inplace=True)
#将Close列添加到最后
close_col = stock.pop('Close')
stock.loc[:,'Close'] = close_col
#填补Dividend缺失值、删除具有缺失值的行
stock["ExpectedDividend"] = stock["ExpectedDividend"].fillna(0)
stock.dropna(inplace=True)
#恢复索引
stock.index = range(stock.shape[0])
二、数据分割与数据重组
# 二、数据分割与数据重组
# 1、数据分割
train_size = int(len(stock) * 0.67)
test_size = len(stock) - train_size
train, test = stock[:train_size], stock[train_size:] # train (6580,12) test(3242,12)
# 2、带标签滑窗
def create_multivariate_dataset_2(dataset, window_size, pred_len): #
"""
将多变量时间序列转变为能够用于训练和预测的数据【带标签的滑窗】
参数:
dataset: DataFrame，其中包含特征和标签，特征从索引3开始，最后一列是标签
window_size: 滑窗的窗口大小
pred_len：多步预测的预测范围/预测步长
"""
X, y, y_indices = [], [], []
for i in range(len(dataset) - window_size - pred_len + 1): # (len-ws-pl+1) --> (6580-30-5+1) = 6546
# 选取从第4列到最后一列的特征和标签
feature_and_label = dataset.iloc[i:i + window_size, 3:].values # (ws,fs_la) --> (30,9)
# 下一个时间点的标签作为目标
target = dataset.iloc[(i + window_size):(i + window_size + pred_len), -1] # pred_len --> 5
# 记录本窗口中要预测的标签的时间点
target_indices = list(range(i + window_size, i + window_size + pred_len)) # pl*(len-ws-pl+1) --> 5*6546 = 32730
X.append(feature_and_label)
y.append(target)
#将每个标签的索引添加到y_indices列表中
y_indices.extend(target_indices)
X = torch.FloatTensor(np.array(X, dtype=np.float32))
y = torch.FloatTensor(np.array(y, dtype=np.float32))
return X, y, y_indices
# 3、数据重组
window_size = 30 #窗口大小
pred_len = 5 #多步预测的步数
X_train_2, y_train_2, y_train_indices = create_multivariate_dataset_2(train, window_size, pred_len) # x(6546,30,9) y(6546,5) (32730,)
X_test_2, y_test_2, y_test_indices = create_multivariate_dataset_2(test, window_size, pred_len) # x(3208,30,9) y(3208,5) (16040,)
三、网络架构与参数设置
# 三、网络架构与参数设置
# 1、定义架构
class MyLSTM(nn.Module):
def __init__(self,input_dim, seq_length, output_size, hidden_size, num_layers):
super().__init__()
self.lstm = nn.LSTM(input_size=input_dim, hidden_size=hidden_size, num_layers=num_layers, batch_first=True)
self.linear = nn.Linear(hidden_size, output_size)
def forward(self, x):
x, _ = self.lstm(x)
#现在我要的是最后一个时间步，而不是全部时间步了
x = self.linear(x[:,-1,:])
return x
# 2、参数设置
input_size = 9 #输入特征的维度
hidden_size = 20 #LSTM隐藏状态的维度
n_epochs = 2000 #迭代epoch
learning_rate = 0.001 #学习率
num_layers = 1 #隐藏层的层数
output_size = 5
#设置GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
# 加载数据，将数据分批次
loader = data.DataLoader(data.TensorDataset(X_train_2, y_train_2), shuffle=True, batch_size=8)
# 3、实例化模型
model = MyLSTM(input_size, window_size, pred_len,hidden_size, num_layers).to(device)
optimizer = optim.Adam(model.parameters(),lr=learning_rate) #定义优化器
loss_fn = nn.MSELoss() #定义损失函数
loader = data.DataLoader(data.TensorDataset(X_train_2, y_train_2)
#每个表单内部是保持时间顺序的即可，表单与表单之间可以shuffle
, shuffle=True
, batch_size=8) #将数据分批次
四、实际训练流程
# 四、实际训练流程
# 初始化早停参数
early_stopping_patience = 3 # 设置容忍的epoch数，即在这么多epoch后如果没有改进就停止
early_stopping_counter = 0 # 用于跟踪没有改进的epoch数
best_train_rmse = float('inf') # 初始化最佳的训练RMSE
train_losses = []
test_losses = []
for epoch in range(n_epochs):
model.train()
for X_batch, y_batch in loader:
y_pred = model(X_batch.to(device))
loss = loss_fn(y_pred, y_batch.to(device))
optimizer.zero_grad()
loss.backward()
optimizer.step()
#验证与打印
if epoch % 10 == 0:
model.eval()
with torch.no_grad():
y_pred = model(X_train_2.to(device)).cpu()
train_rmse = np.sqrt(loss_fn(y_pred, y_train_2))
y_pred = model(X_test_2.to(device)).cpu()
test_rmse = np.sqrt(loss_fn(y_pred, y_test_2))
print("Epoch %d: train RMSE %.4f, test RMSE %.4f" % (epoch, train_rmse, test_rmse))
# 将当前epoch的损失添加到列表中
train_losses.append(train_rmse)
test_losses.append(test_rmse)
# 早停检查
if train_rmse < best_train_rmse:
best_train_rmse = train_rmse
early_stopping_counter = 0 # 重置计数器
else:
early_stopping_counter += 1 # 增加计数器
if early_stopping_counter >= early_stopping_patience:
print(f"Early stopping triggered after epoch {epoch}. Training RMSE did not decrease for {early_stopping_patience} consecutive epochs.")
break # 跳出训练循环
结果显示：
Epoch 0: train RMSE 1470.9308, test RMSE 1692.0652
Epoch 5: train RMSE 1415.7896, test RMSE 1639.1147
Epoch 10: train RMSE 1364.8196, test RMSE 1590.2207
......
Epoch 100: train RMSE 654.3458, test RMSE 904.7958
Epoch 105: train RMSE 638.2536, test RMSE 886.3511
Epoch 110: train RMSE 625.7336, test RMSE 870.9800
......
Epoch 200: train RMSE 598.3364, test RMSE 820.4078
Epoch 205: train RMSE 598.3354, test RMSE 820.3406
Epoch 210: train RMSE 598.3349, test RMSE 820.2874
......
Epoch 260: train RMSE 598.3341, test RMSE 820.1312
Epoch 265: train RMSE 598.3341, test RMSE 820.1294
Early stopping triggered after epoch 265. Training RMSE did not decrease for 3 consecutive epochs.
五、可视化结果
# 五、可视化结果
# 1、损失曲线
plt.figure(figsize=(10, 5))
plt.plot(train_losses, label='Train RMSE')
plt.plot(test_losses, label='Test RMSE')
plt.xlabel('Epochs')
plt.ylabel('RMSE')
plt.title('Train and Test RMSE Over Epochs')
plt.legend()
plt.show()
结果分析：预测效果不是很好，考虑进行数据预处理和特征工程
【扩展】股票数据的数据预处理与特征工程（后续更新~）
