用LSTM模型预测股价的例子（1）

基本说明

本实例用的是单个参数“收盘价”，学习后。用10天的收盘价预测后面1天的收盘价。
数据如下图：

后续我们还要采用这个数据，进一步添加其他的影响因子进行预测。

直接上代码

代码中包含注释，我就不多说了

import tensorflow as tf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


# 1. 加载数据
def load_stock_data(file_path):
    data = pd.read_csv(file_path)
    return data['close'].values


# 2. 划分训练集和测试集
def prepare_data(data, train_ratio=0.8):
    train_size = int(len(data) * train_ratio)
    train_data = data[:train_size]
    test_data = data[train_size:]
    return train_data, test_data


# 3. 创建数据集
def create_dataset(data, time_steps):
    Xs, ys = [], []
    for i in range(len(data) - time_steps):
        v = data[i:(i + time_steps)]
        Xs.append(v)
        ys.append(data[i + time_steps])
    return np.array(Xs), np.array(ys)


# 4. 定义 LSTM 模型
def build_model(input_shape):
    model = tf.keras.Sequential([
        tf.keras.layers.LSTM(units=50, return_sequences=True, input_shape=input_shape),
        tf.keras.layers.LSTM(units=50),
        tf.keras.layers.Dense(units=1)
    ])
    model.compile(optimizer='adam', loss='mse')
    return model


# 5. 训练模型
def train_model(model, X_train, y_train, epochs=20):
    history = model.fit(
        X_train, y_train,
        epochs=epochs,
        batch_size=32,
        validation_split=0.1,
        shuffle=False
    )
    return history


# 6. 预测和可视化结果
def predict_and_visualize(model, X_train, y_train, X_test, y_test):
    train_predict = model.predict(X_train)
    test_predict = model.predict(X_test)

    plt.figure(figsize=(10, 6))
    plt.plot(y_train, label='True Train')
    plt.plot(train_predict, label='Predicted Train')
    plt.plot(range(len(y_train), len(y_train) + len(y_test)), y_test, label='True Test')
    plt.plot(range(len(y_train), len(y_train) + len(y_test)), test_predict, label='Predicted Test')
    plt.legend(loc='upper left')
    plt.show()


if __name__ == "__main__":
    file_path ='d:/test.csv'
    data = load_stock_data(file_path)
    train_data, test_data = prepare_data(data)
    time_steps = 10
    X_train, y_train = create_dataset(train_data, time_steps)
    X_test, y_test = create_dataset(test_data, time_steps)

    # 数据归一化
    from sklearn.preprocessing import MinMaxScaler
    scaler = MinMaxScaler(feature_range=(0, 1))
    X_train = scaler.fit_transform(X_train.reshape(-1, 1)).reshape(X_train.shape)
    y_train = scaler.transform(y_train.reshape(-1, 1)).reshape(y_train.shape)
    X_test = scaler.transform(X_test.reshape(-1, 1)).reshape(X_test.shape)
    y_test = scaler.transform(y_test.reshape(-1, 1)).reshape(y_test.shape)

    #print(X_train.shape)
    input_shape = (X_train.shape[1], 1)
    model = build_model(input_shape)
    history = train_model(model, X_train, y_train, epochs=20)
    predict_and_visualize(model, X_train, y_train, X_test, y_test)

预测结果

红色的是预测的，绿色的是实际的。