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') # 初始化最佳的训练RMSEtrain_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_rmseearly_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()
结果分析：预测效果不是很好，考虑进行数据预处理和特征工程
【扩展】股票数据的数据预处理与特征工程（后续更新~）
