python, LSTM을 활용한 주식 가격 예측

오늘은 python을 이용하여 과거의 종가 기준으로 주가를 예측하는 코드를 만들어 보았습니다.

먼저 LSTM이 무엇인지는 자세히 기술하지 않겠습니다. LSTM(Long Short Time Memory)의 약자로 이전 정보를 오랫동안 기억할 수 있는 메모리 셀을 가지고 있으며, 이를 통해 긴 시퀀스 데이터를 처리할 수 있는 모델입니다.

검색해보면 시계열 데이터의 예측, 자연어 처리, 음성 인식, 이미지 분류 등에서 중요한 역할을 하는 모델 중 하나라고 합니다. 

 

1. Data 가져오기

이제 본론으로 들어가서, 주식가격을 예측하려면 이전의 데이터들을 가지고 학습을 시켜야 합니다. 학습을 시키려면 이전의 데이터들이 필요한데, 이는 yfinance에서 다운로드 받을 수 있습니다.

import yfinance as yf

# 삼성전자의 주식 심볼
symbol = '005930.KS'

# 데이터를 가져올 날짜 범위
start_date = '2015-01-01'
end_date = '2024-08-29'

data = yf.download(symbol, start=start_date, end=end_date)

위와 같이 하면 삼성전자의 주가데이터를 가져오게 됩니다.

 

2. Data 분할, 정규화

이건 뒤에 LSTM 모델에 적용하려고 Data를 나누고, 그 형태를 만든다는 의미입니다.

from sklearn.preprocessing import MinMaxScaler

# 종가 정보만을 포함하는 새로운 데이터프레임을 생성합니다
closing_prices = data['Close'].to_frame()

# 데이터 분할
train_data = closing_prices['2015-01-01':'2022-12-31']
test_data = closing_prices['2023-01-01':'2024-08-29']

# 데이터 정규화
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_train_data = scaler.fit_transform(train_data)
scaled_test_data = scaler.transform(test_data)

주석과 같이 앞에서 가져온 정보에서 종가 정보만 취합니다.

그리고 종가정보를 또 둘로 나눕니다. 학습은 2015년1월1일 ~ 2022년12월31일까지의 종가로 학습시키고, 2023년1월1일부터 2024년08월29일까지는 학습한 정보를 가지고 어떻게 될 지 테스트 해보는거에요. 이건 실제 삼성전자의 주가가 아니라 앞의 학습한 데이터로 예측한다는 의미에요. 이거를 실제 삼성전자의 주가와 비교해 보면 이 LSTM모델의 믿을만 한지 볼 수 있겠죠?

 

3. Dataset 생성

# 데이터셋 생성 함수
def create_dataset(data, time_step=1):
    X, Y = [], []
    for i in range(len(data) - time_step - 1):
        a = data[i:(i + time_step), 0]
        X.append(a)
        Y.append(data[i + time_step, 0])
    return np.array(X), np.array(Y)

# 시퀀스 길이 설정
time_step = 60

# 학습 데이터셋 생성
X_train, y_train = create_dataset(scaled_train_data, time_step)
X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)

# 테스트 데이터셋 생성
X_test, y_test = create_dataset(scaled_test_data, time_step)
X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)

앞의 데이터를 이제 LSTM에 넣을 Parameter 형태로 만드는 작업입니다.

 

4. LSTM 모델 생성

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense

# LSTM 모델 생성
model = Sequential()
model.add(LSTM(50, return_sequences=True, input_shape=(time_step, 1)))
model.add(LSTM(50, return_sequences=False))
model.add(Dense(25))
model.add(Dense(1))

# 모델 컴파일
model.compile(optimizer='adam', loss='mean_squared_error')

# 모델 학습
model.fit(X_train, y_train, batch_size=1, epochs=1)

# 예측
train_predict = model.predict(X_train)
test_predict = model.predict(X_test)

# 데이터 역정규화
train_predict = scaler.inverse_transform(train_predict)
test_predict = scaler.inverse_transform(test_predict)

이제 python에서 LSTM 모델을 사용하려면 위와 같이 생성하고 위에서 만들어둔 dataset을 넣으면 됩니다.

LSTM모델 생성시에 들어가는 숫자들은 각각 의미가 있는데요, 여기서 주로 의미있는 값은 time_step과 epochs 값입니다.

time_step은  LSTM에 입력되는 시퀀스 데이터의 시간 단계 또는 길이를 의미합니다. 예를 들어, 주식 가격 예측에서 timestep이 10이라면, 모델은 이전 10일의 데이터를 사용하여 다음 날의 가격을 예측합니다.

epochs는 전체 훈련 데이터셋이 모델을 통해 한 번 통과하는 과정을 의미합니다. 즉, epochs가 50이라면, 모델은 훈련 데이터셋을 50번 반복하여 학습하게 됩니다.

이 두 가지 값은 LSTM 모델의 학습 과정에서 매우 중요한 역할을 하므로, 실험을 통해 최적의 값을 찾는 것이 필요합니다. 정답이 없는 값이에요. 

 

5. 결과 시각화

이제 위에서 학습을 통해 얻은 결과를 눈으로 봐야겠죠?

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

# 결과 시각화
plt.figure(figsize=(14, 5))
plt.plot(closing_prices.index, closing_prices['Close'], label='Actual Price')
plt.plot(train_data.index[time_step+1:], train_predict, label='Train Predict')
plt.plot(test_data.index[time_step+1:], test_predict, label='Test Predict')

# 마지막 날짜 이후 60일 동안의 예측
last_60_days = closing_prices[-time_step:].values
last_60_days_scaled = scaler.transform(last_60_days)

X_last = []
X_last.append(last_60_days_scaled)
X_last = np.array(X_last)
X_last = np.reshape(X_last, (X_last.shape[0], X_last.shape[1], 1))

pred_price = []
for _ in range(60):
    pred = model.predict(X_last)
    pred_price.append(pred[0])
    X_last = np.append(X_last[:,1:,:], [pred], axis=1)

pred_price = scaler.inverse_transform(pred_price)

prediction_dates = pd.date_range(start=closing_prices.index[-1], periods=61)[1:]
plt.plot(prediction_dates, pred_price, label='Predicted Price for the next 60 days')

plt.xlabel('Date')
plt.ylabel('Close Price')
plt.legend()
plt.show()

저는 여기에 마지막 날짜 이후 60일동안의 예측 값도 구해서 표시하도록 해 보았습니다.

이제 위의 모든 코드들을 합쳐서 돌려보면 다음과 같이 화면에 표시됩니다.

학습중

결과화면

위와 같이 실제가격정보, 학습하여 예측한 결과들이 표시됩니다. 위에서 언급했던 time_step과 epochs 값을 다르게 하여 돌릴 때마다 예측 그래프가 달라집니다.

 

그대로 복사하시기 편하도록 풀코드는 아래와 같습니다.

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from datetime import timedelta
import yfinance as yf

# 삼성전자의 주식 심볼
symbol = '005930.KS'

# 데이터를 가져올 날짜 범위
start_date = '2015-01-01'
end_date = '2024-08-29'

# yfinance를 이용하여 주식 데이터를 가져옵니다
data = yf.download(symbol, start=start_date, end=end_date)

# 종가 정보만을 포함하는 새로운 데이터프레임을 생성합니다
closing_prices = data['Close'].to_frame()

# 데이터 분할
train_data = closing_prices['2015-01-01':'2022-12-31']
test_data = closing_prices['2023-01-01':'2024-08-29']

# 데이터 정규화
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_train_data = scaler.fit_transform(train_data)
scaled_test_data = scaler.transform(test_data)

# 데이터셋 생성 함수
def create_dataset(data, time_step=1):
    X, Y = [], []
    for i in range(len(data) - time_step - 1):
        a = data[i:(i + time_step), 0]
        X.append(a)
        Y.append(data[i + time_step, 0])
    return np.array(X), np.array(Y)

# 시퀀스 길이 설정
time_step = 60

# 학습 데이터셋 생성
X_train, y_train = create_dataset(scaled_train_data, time_step)
X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)

# 테스트 데이터셋 생성
X_test, y_test = create_dataset(scaled_test_data, time_step)
X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)

# LSTM 모델 생성
model = Sequential()
model.add(LSTM(50, return_sequences=True, input_shape=(time_step, 1)))
model.add(LSTM(50, return_sequences=False))
model.add(Dense(25))
model.add(Dense(1))

# 모델 컴파일
model.compile(optimizer='adam', loss='mean_squared_error')

# 모델 학습
model.fit(X_train, y_train, batch_size=1, epochs=1)

# 예측
train_predict = model.predict(X_train)
test_predict = model.predict(X_test)

# 데이터 역정규화
train_predict = scaler.inverse_transform(train_predict)
test_predict = scaler.inverse_transform(test_predict)

# 결과 시각화
plt.figure(figsize=(14, 5))
plt.plot(closing_prices.index, closing_prices['Close'], label='Actual Price')
plt.plot(train_data.index[time_step+1:], train_predict, label='Train Predict')
plt.plot(test_data.index[time_step+1:], test_predict, label='Test Predict')

# 마지막 날짜 이후 60일 동안의 예측
last_60_days = closing_prices[-time_step:].values
last_60_days_scaled = scaler.transform(last_60_days)

X_last = []
X_last.append(last_60_days_scaled)
X_last = np.array(X_last)
X_last = np.reshape(X_last, (X_last.shape[0], X_last.shape[1], 1))

pred_price = []
for _ in range(60):
    pred = model.predict(X_last)
    pred_price.append(pred[0])
    X_last = np.append(X_last[:,1:,:], [pred], axis=1)

pred_price = scaler.inverse_transform(pred_price)

prediction_dates = pd.date_range(start=closing_prices.index[-1], periods=61)[1:]
plt.plot(prediction_dates, pred_price, label='Predicted Price for the next 60 days')

plt.xlabel('Date')
plt.ylabel('Close Price')
plt.legend()
plt.show()

 

6. 마치며...

이 코드는 이런게 있구나.. 하는 느낌으로 봐 주세요. 이 결과를 믿고 그대로 투자하시는 분은 없겠죠?

다음에도 재미있는 주제로 글 올리겠습니다. 좋은 하루 되세요.