ModelTrainingPython/TVS_DL/LSTM_Test.py

# 导入相关库
import os

import pandas as pd
import torch  # 导入PyTorch库
import torch.nn as nn  # 导入神经网络模块
import tushare as ts  # 导入tushare库，用于获取股票数据
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类，用于加载数据


# 获取数据
class GetData:
    def __init__(self, stock_id, save_path):
        """
        初始化方法
        :param stock_id: 股票id
        :param save_path: 数据保存路径
        """
        self.min_value = None
        self.max_value = None
        self.stock_id = stock_id
        self.save_path = save_path
        self.data = None

    def 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.data

    def 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模型: 单层单向LSTM网络+全连接层输出
class Model(nn.Module):
    def __init__(self, n):
        # 初始化方法
        super(Model, self).__init__()  # 调用父类的初始化方法
        # 定义LSTM层 输入大小为n, 隐藏层大小为256,批次优先为True
        self.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_size
        batch_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_num

        for 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

        # 如果当前测试损失小于最佳损失，则更新最佳模型和轮次计数器 否则 轮次计数器加1
        if test_loss < best_loss:
            best_loss = test_loss
            best_model = copy(model)  # ✅ 使用 copy(model)，而不是 copy(Model
            torch.save(best_model.state_dict(), './best_model.pth')
            epoch_cnt = 0
        else:
            epoch_cnt += 1
        if epoch_cnt > early_stop:
            break


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 = 0
    total_test_num = 0

    for 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_num
    return 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()


if __name__ == '__main__':
    # 超参数
    days_num = 5  # 天数
    epoch = 100  # 训练轮次
    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}")