LSTM_1.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. import pandas as pd
  2. import numpy as np
  3. from sklearn.preprocessing import MinMaxScaler
  4. from keras.models import Sequential
  5. from keras.layers import Dense, LSTM
  6. from ta.volatility import BollingerBands
  7. from ta.trend import MACD
  8. from ta.trend import EMAIndicator
  9. from ta.volume import MFIIndicator
  10. from sqlalchemy import create_engine
  11. # 加载数据
  12. engine_tech = create_engine('mysql+pymysql://root:r6kEwqWU9!v3@localhost:3307/qmt_stocks_tech?charset=utf8')
  13. df = pd.read_sql_table('600000.SH_1d', con=engine_tech)
  14. df = df[['open_back', 'high_back', 'low_back', 'close_back', 'volume_back', 'close_back', 'macd', 'willr', 'volume_back', 'dif', 'dea', 'rsi_6', 'rsi_12', 'rsi_24']].values
  15. # df = df[~np.isnan(df).any(axis=1), :]
  16. # 归一化
  17. scaler = MinMaxScaler(feature_range=(0, 1))
  18. scaled_data = scaler.fit_transform(df)
  19. print(len(df), len(scaled_data))
  20. # 构造训练数据
  21. lookback = 90
  22. target_days = 30
  23. x_train = []
  24. y_train = []
  25. for i in range(lookback, len(df) - target_days):
  26. x_train.append(scaled_data[i - lookback:i])
  27. y_train.append(scaled_data[i:i + target_days, 0])
  28. x_train, y_train = np.array(x_train), np.array(y_train)
  29. print('11111111111111111111', x_train.shape, y_train.shape)
  30. # 构建模型
  31. model = Sequential()
  32. model.add(LSTM(units=50, return_sequences=True, input_shape=(x_train.shape[1], x_train.shape[2])))
  33. model.add(LSTM(units=50, return_sequences=True))
  34. model.add(LSTM(units=50))
  35. model.add(Dense(units=30))
  36. model.compile(optimizer='adam', loss='mean_squared_error')
  37. model.fit(x_train, y_train, epochs=10, batch_size=32)
  38. # 预测未来30天走势
  39. test_data = scaled_data[-lookback:]
  40. test_data = np.expand_dims(test_data, axis=0)
  41. print(test_data.shape)
  42. predictions = model.predict(test_data)
  43. # 将归一化的预测数据转换为原始数据
  44. predictions = scaler.inverse_transform(predictions)
  45. # 打印预测结果
  46. print(predictions)