当前位置：首页 > news >正文

PyTorch LSTM练习案例：股票成交量趋势预测

news 2025/7/5 1:48:06

文章目录

案例介绍
源码地址
代码实现
- 导入相关库
- 数据获取和处理
- 搭建LSTM模型
- 训练模型
- 测试模型
- 绘制折线图
- 主函数
绘制结果

案例介绍

本例使用长短期记忆网络模型对上海证券交易所工商银行的股票成交量做一个趋势预测，这样可以更好地掌握股票买卖点，从而提高自己的收益率。

源码地址

stock_prediction

代码实现

导入相关库

import os
import pandas as pd
import torch  # 导入PyTorch库
import torch.nn as nn  # 导入神经网络模块
import torch.optim as optim  # 导入优化器模块
from tqdm import tqdm  # 导入tqdm库，用于显示进度条
import matplotlib.pyplot as plt  # 导入matplotlib库，用于绘制图表
from copy import deepcopy as copy  # 导入deepcopy函数，用于深拷贝对象
from torch.utils.data import DataLoader, TensorDataset  # 导入DataLoader和TensorDataset类，用于加载数据

数据获取和处理

data.csv文件中包括开盘价，收盘价，最高价，最低价，成交量5个数据特征，使用每天的收盘价作为学习目标，每个样本都包含连续几天的数据作为一个序列样本，然后，将数据划分为训练集和测试集供后续使用。
GetData类用于获取和处理数据，它具有以下方法：
1. init(self, stock_id,save_path)：初始化方法，接受股票ID和数据保存路径作为参数，并将它们存储在实例变量中。
2. getData(self)：获取数据的方法，获取股票历史数据并进行处理，然后保存到文件中。返回处理后的数据。
3. process_data(self,n)：处理数据的方法,将数据分为特征和标签，并划分为训练集和测试集。接受滑动窗口大小n作为参数。如果数据为空，则调用getData方法获取数据。返回训练集的特征、测试集的特征、训练集的标签和测试集的标签。

# 获取数据
class GetData:def __init__(self, stock_id, save_path):"""初始化方法:param stock_id: 股票id:param save_path: 数据保存路径"""self.min_value = Noneself.max_value = Noneself.stock_id = stock_idself.save_path = save_pathself.data = Nonedef getData(self):"""获取数据数据处理并保存:return: None""""""# 获取股票数据self.data = ts.get_hist_data(self.stock_id).iloc[::-1]# 选取特定列作为数据self.data = self.data[["open", "close", "high", "low", "volume"]]# 计算数据列的最大值和最小值self.max_value = self.data['volume'].max()self.min_value = self.data['volume'].min()# 归一化处理self.data = self.data.apply(lambda x: (x - min(x)) / (max(x) - min(x)))# 保存数据self.data.to_csv(self.save_path)return self.data"""# 本地数据data.csv由于是归一化后的数据，所以最大值和最小值并不准确，所以运行结果会有误差，重在体验整个项目的逻辑即可columns = ['open', 'close', 'high', 'low', 'volume']self.data = pd.read_csv(self.save_path, names=columns, header=0)# 计算数据列的最大值和最小值self.max_value = self.data['volume'].max()self.min_value = self.data['volume'].min()return self.datadef process_data(self, n):"""处理数据:param n: 滑动窗口大小:return: 训练集的特征、测试集的特征、训练集的标签、测试集的标签"""if self.data is None:self.getData()# 提取特征和标签数据"""iloc 是 Pandas 库中用于按位置索引选取数据的方法"""feature = [self.data.iloc[i: i + n].values.tolist()for i in range(len(self.data) - n + 2)if i + n < len(self.data)]label = [self.data.close.values[i + n]for i in range(len(self.data) - n + 2)if i + n < len(self.data)]# 划分训练集和数据集train_x = feature[:500]test_x = feature[500:]train_y = label[:500]test_y = label[500:]return train_x, test_x, train_y, test_y

搭建LSTM模型

定义一个Model的神经网络模块，包含一个LSTM层和线性层，在初始化方法中接受一个参数n，并创建一个LSTM层和一个线性层。
在前向传播方法中，通过LSTM层处理输入x得到输出lstm_output和隐藏状态hidden_state, cell_state，然后通过线性处理num_layers得到最终输出final_output ，最终返回final_output `作为模型的输出。

# 搭建LSTM模型: 单层单向LSTM网络+全连接层输出
class Model(nn.Module):def __init__(self, n):# 初始化方法super(Model, self).__init__()  # 调用父类的初始化方法# 定义LSTM层 输入大小为n, 隐藏层大小为256,批次优先为Trueself.lstm_layer = nn.LSTM(input_size=n, hidden_size=256, batch_first=True)#  定义全连接层 输入特征数为256, 输出特征数为1 有偏差self.linear_layer = nn.Linear(in_features=256, out_features=1, bias=True)# 向前传播方法def forward(self, x):"""x: 输入数据（通常是时间序列的特征）lstm_output: LSTM 层的输出序列hidden_state: LSTM 的隐藏状态（用于传递长期记忆）cell_state: LSTM 的细胞状态（仅在 LSTM 中存在）final_output: 经过全连接层后的最终输出"""# LSTM 层的前向传播，得到输出和隐藏状态lstm_output, (hidden_state, cell_state) = self.lstm_layer(x)# 获取隐藏状态的维度：batch_size, num_layers, hidden_sizebatch_size, num_layers, hidden_size = hidden_state.shape# 将隐藏状态输入全连接层，需要先展平为二维final_output = self.linear_layer(hidden_state.view(batch_size * num_layers, hidden_size))return final_output

训练模型

模型训练包括训练，测试，损失计算和模型保存等功能。

# 训练模型
def train_model(epoch, train_dataloader, test_dataloader, optimizer, early_stop, model):"""训练模型的函数:param model: 模型:param early_stop: 提前停止的轮数:param optimizer: 优化器:param epoch: 训练轮次:param train_dataloader: 训练数据加载器:param test_dataloader: 测试数据加载器:return:"""best_model = None  # 用于保存最佳模型train_loss = 0  # 训练损失test_loss = 0  # 测试损失best_loss = 100  # 最佳损失epoch_cnt = 0  # 训练轮次计数器for i in range(epoch):total_train_loss = 0  # 训练总损失total_train_num = 0  # 训练总样本数total_test_loss = 0  # 测试总损失total_test_num = 0  # 测试总样本数for x, y in tqdm(train_dataloader, desc=f"Epoch:{i} | Train Loss:{train_loss} | Test Loss:{test_loss}"):x_num = len(x)  # 当前批次样本数p = model(x)  # 模型预测   ✅ 使用 model(x)，而不是 Model(x)loss = loss_func(p, y)  # 计算损失optimizer.zero_grad()  # 清空梯度loss.backward()  # 反向传播optimizer.step()  # 更新参数total_train_loss += loss.item()  # 训练损失累加total_train_num += x_num  # 训练样本数累加#  计算训练损失train_loss = total_train_loss / total_train_numfor x, y in test_dataloader:x_num = len(x)  # 当前批次样本数p = model(x)  # 模型预测  ✅ 使用 model(x)，而不是 Model(x)loss = loss_func(p, y)  # 计算损失optimizer.zero_grad()  # 清空梯度loss.backward()  # 反向传播optimizer.step()  # 更新参数total_test_loss += loss.item()  # 测试损失累加total_test_num += x_num  # 测试样本数累加test_loss = total_test_loss / total_test_num# 如果当前测试损失小于最佳损失，则更新最佳模型和轮次计数器 否则 轮次计数器加1if test_loss < best_loss:best_loss = test_lossbest_model = copy(model)  # ✅ 使用 copy(model)，而不是 copy(Modeltorch.save(best_model.state_dict(), './best_model.pth')epoch_cnt = 0else:epoch_cnt += 1if epoch_cnt > early_stop:break

测试模型

在代码中定义一个名为test_model的函数，用于测试模型并返回预测值、真实标签以及测试损失，函数接收一个名为test_dataLoader_的DataLoader参数，其中包含测试数据。
在函数内部，首先创建了空的预测值列表pred和真实标签列表label。然后创建了一个模型对象model_f，加载了预先保存的模型状态字典./best_model.pth，并将模型设置为评估模式。
接着，通过遍历test_dataLoader中的数据进行预测。对每个数据样本x，模型预测p，计算损失
值并累加到total_test_loss中。同时，将预测值和真实标签分别添加到pred和label列表中。
最后，计算平均测试损失test_loss，并将预测值列表pred、真实标签列表label和测试损失test_loss作为结果返回。

def test_model(test_dataloader):"""测试模型，并返回预测值、真实标签和测试损失:param test_dataloader: 测试数据加载器:return: pred,label,test_loss"""pred = []  # 预测值列表label = []  # 真实标签列表model_f = Model(5)  # 创建模型对象model_f.load_state_dict(torch.load('./best_model.pth'))  # 加载最佳模型model_f.eval()  # 设置模型为评估模式total_test_loss = 0total_test_num = 0for x, y in test_dataloader:x_num = len(x)p = model_f(x)  # ✅ 使用 model_f(x)loss = loss_func(p, y)total_test_loss += loss.item()total_test_num += x_num# 将预测值和真实标签添加到列表中pred.extend(p.data.squeeze(1).tolist())label.extend(y.data.tolist())# 获取预测值和真实标签test_loss = total_test_loss / total_test_numreturn pred, label, test_loss

绘制折线图

绘制股票日成交量的折线图，并输出模型测试集的损失。

def plot_img(data, pred):"""绘制真实值与预测值对比图:param data: 真实标签列表:param pred: 模型预测值列表:return:"""# 设置支持中文的字体plt.rcParams['font.sans-serif'] = ['SimHei']plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题plt.figure(figsize=(14, 8))# 绘制真实值曲线plt.plot(range(len(data)), data, color='blue', label='真实值（收盘价）', linewidth=2)# 绘制预测值曲线plt.plot(range(len(pred)), pred, color='green', label='预测值（模型输出）', linestyle='--', linewidth=2)# 添加预测区间（每5个点绘制一个3天的预测区间）for i in range(0, len(pred) - 3, 5):price = [data[i] + pred[j] - pred[i] for j in range(i, i + 3)]plt.plot(range(i, i + 3), price, color='red', alpha=0.6, linestyle=':', linewidth=1.5)# 设置标题和标签plt.title('股票价格预测结果对比', fontsize=20)plt.xlabel('时间步（天数）', fontsize=16)plt.ylabel('股票收盘价(亿)', fontsize=16)# 设置刻度字体plt.xticks(fontproperties='Times New Roman', size=14)plt.yticks(fontproperties='Times New Roman', size=14)# 显示图例plt.legend(loc='upper left', fontsize=14)# 显示网格plt.grid(True, linestyle='--', alpha=0.5)# 展示图形plt.tight_layout()plt.show()

主函数

初始化设置
- 定义超参数（训练轮次、批次大小等）。
- 创建模型和数据加载对象，指定股票ID和数据保存路径。
数据处理
- 获取股票数据，按时间窗口（5天）生成输入序列和标签。
- 将数据转换为PyTorch张量，并分批次加载（batch_size=20）。
模型训练
- 使用均方误差损失和Adam优化器训练模型。
- 监控验证损失，若连续5轮无改进则提前停止，保存最佳模型。
模型测试
- 加载保存的最佳模型，在测试集上预测并计算损失。
- 将预测值和真实值反归一化，还原为原始价格。
结果输出
- 绘制预测值与真实值的对比图。
- 打印测试集上的最终损失值。

设置参数 → 加载并预处理数据 → 训练模型（含早停） → 测试并还原预测结果 → 可视化输出。

if __name__ == '__main__':# 超参数days_num = 5  # 天数epoch = 20  # 训练轮次fea = 5  # 特征数量batch_size = 20  # 批次大小early_stop = 5  # 提前停止轮次# 创建模型对象model = Model(fea)#  创建数据加载器gd = GetData(stock_id='601398', save_path='./data.csv')train_x, test_x, train_y, test_y = gd.process_data(days_num)# 将数据转换为张量train_x = torch.tensor(train_x).float()test_x = torch.tensor(test_x).float()train_y = torch.tensor(train_y).float()test_y = torch.tensor(test_y).float()# 构建训练数据集和测试数据集train_data = TensorDataset(train_x, train_y)train_dataloader = DataLoader(train_data, batch_size=batch_size)test_data = TensorDataset(test_x, test_y)test_dataloader = DataLoader(test_data, batch_size=batch_size)# 创建损失函数和优化器loss_func = nn.MSELoss()optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练模型train_model(epoch, train_dataloader, test_dataloader, optimizer, early_stop, model)# 只有模型存在时才进行测试if os.path.exists('./best_model.pth'):pred, label, test_loss = test_model(test_dataloader)else:print("模型文件不存在，请先完成训练并确保模型已保存。")#  将预测值和真实标签转换为真实价格pred = [ele * (gd.max_value - gd.min_value) + gd.min_value for ele in pred]data = [ele * (gd.max_value - gd.min_value) + gd.min_value for ele in label]# 绘制图像plot_img(data, pred)print(f"模型损失：{test_loss}")