파이썬 AI 주가 예측 모델 만들기 A to Z (ft. 텐서플로우, LSTM)

pip install tensorflow pandas numpy matplotlib yfinance scikit-learn

import yfinance as yf
import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt

# 1. 데이터 불러오기 (삼성전자, 2020-01-01 ~ 2023-12-31)
ticker = "005930.KS"
data = yf.download(ticker, start="2020-01-01", end="2023-12-31")

# 종가(Close) 데이터만 사용
close_data = data[['Close']].values

print(data.head())

# 2. 데이터 정규화 (0과 1 사이의 값으로 변환)
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(close_data)

# 3. 훈련 데이터와 테스트 데이터 분리 (80% 훈련, 20% 테스트)
train_size = int(len(scaled_data) * 0.8)
train_data = scaled_data[:train_size]
test_data = scaled_data[train_size:]

# 4. 시계열 데이터셋 생성 함수
def create_dataset(dataset, time_step=60):
    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 # 과거 60일의 데이터를 기반으로 다음 날을 예측
X_train, y_train = create_dataset(train_data, time_step)
X_test, y_test = create_dataset(test_data, time_step)

# LSTM 모델에 입력하기 위해 데이터 형태 변환 (samples, time_steps, features)
X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)
X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)

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

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

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

model.summary()

# 7. 모델 학습
history = model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=100, batch_size=64, verbose=1)

# 학습 손실 시각화
plt.plot(history.history['loss'], label='train loss')
plt.plot(history.history['val_loss'], label='validation loss')
plt.title('Model Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

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

# 9. 예측 결과 정규화 되돌리기
train_predict = scaler.inverse_transform(train_predict)
y_train_orig = scaler.inverse_transform(y_train.reshape(-1, 1))
test_predict = scaler.inverse_transform(test_predict)
y_test_orig = scaler.inverse_transform(y_test.reshape(-1, 1))


# 10. 예측 결과 시각화
plt.figure(figsize=(12,6))
plt.plot(data.index, data['Close'], label='Actual Price')

# 훈련 데이터 예측 결과 플롯
train_predict_plot = np.empty_like(close_data)
train_predict_plot[:, :] = np.nan
train_predict_plot[time_step:len(train_predict)+time_step, :] = train_predict
plt.plot(data.index, train_predict_plot, label='Train Predict')

# 테스트 데이터 예측 결과 플롯
test_predict_plot = np.empty_like(close_data)
test_predict_plot[:, :] = np.nan
test_predict_plot[len(train_predict)+(time_step*2)+1:len(close_data)-1, :] = test_predict
plt.plot(data.index, test_predict_plot, label='Test Predict')

plt.title('Stock Price Prediction')
plt.xlabel('Date')
plt.ylabel('Price')
plt.legend()
plt.show()

# 11. 간단한 투자 백테스팅
test_start_index = len(train_data) + time_step + 1
actual_prices = scaler.inverse_transform(test_data[time_step+1:])
predicted_prices = test_predict

# 투자 시뮬레이션
initial_capital = 10000000 # 초기 자본금 1천만원
capital = initial_capital
shares = 0
portfolio_values = []

for i in range(len(predicted_prices) - 1):
    # 다음 날 가격이 오를 것으로 예측되면
    if predicted_prices[i+1] > actual_prices[i]:
        # 보유 현금으로 살 수 있는 만큼 최대한 매수
        if capital > 0:
            shares_to_buy = capital // actual_prices[i]
            shares += shares_to_buy
            capital -= shares_to_buy * actual_prices[i]
    # 다음 날 가격이 내릴 것으로 예측되면
    elif predicted_prices[i+1] < actual_prices[i]:
        # 보유 주식 전량 매도
        if shares > 0:
            capital += shares * actual_prices[i]
            shares = 0
    
    # 현재 포트폴리오 가치 계산
    current_value = capital + shares * actual_prices[i]
    portfolio_values.append(current_value)

# 누적 수익률 계산 및 시각화
portfolio_df = pd.DataFrame(portfolio_values, columns=['Portfolio Value'])
portfolio_df['Daily Return'] = portfolio_df['Portfolio Value'].pct_change()
portfolio_df['Cumulative Return'] = (1 + portfolio_df['Daily Return']).cumprod() - 1

plt.figure(figsize=(10, 5))
plt.plot(portfolio_df['Cumulative Return'] * 100)
plt.title('AI Trading Strategy - Cumulative Return')
plt.xlabel('Days')
plt.ylabel('Return (%)')
plt.grid(True)
plt.show()

print(f"Final Portfolio Value: {portfolio_values[-1]:,.0f} KRW")
print(f"Total Return: {((portfolio_values[-1] / initial_capital) - 1) * 100:.2f}%")