基于lstm写一个股价预测代码
时间: 2024-01-25 12:04:59 浏览: 20
好的,以下是一个简单的基于LSTM的股价预测代码,使用Keras库实现:
```python
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential
from keras.layers import Dense, LSTM
# 读取数据
df = pd.read_csv('stock_price.csv')
# 取出收盘价列
data = df.filter(['Close']).values
# 划分训练集和测试集
train_data_len = int(np.ceil(len(data) * 0.8))
train_data = data[0:train_data_len, :]
test_data = data[train_data_len - 60:, :]
# 数据归一化
scaler = MinMaxScaler(feature_range=(0, 1))
train_data = scaler.fit_transform(train_data)
test_data = scaler.transform(test_data)
# 创建训练集和测试集
def create_dataset(dataset, time_step=1):
X, Y = [], []
for i in range(len(dataset) - time_step - 1):
a = dataset[i : (i + time_step), 0]
X.append(a)
Y.append(dataset[i + time_step, 0])
return np.array(X), np.array(Y)
time_step = 60
X_train, Y_train = create_dataset(train_data, time_step)
X_test, Y_test = create_dataset(test_data, time_step)
# 调整输入数据维度
X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], 1))
X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], 1))
# 创建LSTM模型
model = Sequential()
model.add(LSTM(units=50, return_sequences=True, input_shape=(X_train.shape[1], 1)))
model.add(LSTM(units=50))
model.add(Dense(units=1))
# 编译模型
model.compile(optimizer='adam', loss='mean_squared_error')
# 拟合模型
model.fit(X_train, Y_train, validation_data=(X_test, Y_test), epochs=100, batch_size=64, verbose=1)
# 使用模型进行预测
train_predict = model.predict(X_train)
test_predict = model.predict(X_test)
# 反归一化预测结果
train_predict = scaler.inverse_transform(train_predict)
Y_train = scaler.inverse_transform([Y_train])
test_predict = scaler.inverse_transform(test_predict)
Y_test = scaler.inverse_transform([Y_test])
# 计算RMSE
train_rmse = np.sqrt(np.mean(np.power((Y_train - train_predict), 2)))
test_rmse = np.sqrt(np.mean(np.power((Y_test - test_predict), 2)))
print('Train RMSE:', train_rmse)
print('Test RMSE:', test_rmse)
# 绘制预测结果
train_predict_plot = np.empty_like(data)
train_predict_plot[:, :] = np.nan
train_predict_plot[time_step : (len(train_predict) + time_step), :] = train_predict
test_predict_plot = np.empty_like(data)
test_predict_plot[:, :] = np.nan
test_predict_plot[(len(data) - len(test_predict)):, :] = test_predict
plt.plot(scaler.inverse_transform(data))
plt.plot(train_predict_plot)
plt.plot(test_predict_plot)
plt.show()
```
代码中使用的数据是一个包含股票收盘价的CSV文件,可以根据实际情况修改文件路径。代码中使用了80%的数据作为训练集,剩余的20%作为测试集。在训练过程中,使用了LSTM模型,并使用RMSE作为评价指标。最后,将训练和测试的预测结果绘制在原始数据图上。