파이토치(Pytorch): LSTM 기반 주식 예측 모델 구현

• 수업 내용 리마인드 및 아카이빙 목적의 업로드

---

---

 

이번 글에서는 LSTM(Long Short-Term Memory) 모델을 사용해 주식 데이터를 예측하는 과정을 단계별로 설명하려고 합니다. 주식 시장의 데이터는 시간에 따라 변화하는 시계열 데이터로, 과거 데이터를 바탕으로 미래의 주가를 예측하는 데 LSTM은 RNN의 한계를 보완한 적합한 모델입니다.

 

1. 필요한 모듈 임포트

먼저, 필요한 모듈을 모두 임포트하고 GPU가 사용 가능한 경우 이를 설정합니다.

import torch
import torch.nn as nn
import torch.optim as optim
from torchinfo import summary
from torch.utils.data import Dataset, DataLoader
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, r2_score
from tqdm import tqdm
import warnings

# 경고 무시
warnings.simplefilter("ignore")

# 시드 설정
torch.manual_seed(10)
torch.cuda.manual_seed(10)

# GPU가 사용 가능하면 GPU, 아니면 CPU 사용
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

 

2. 데이터 불러오기

이번 프로젝트에서는 Yahoo Finance에서 제공하는 AAPL(Apple Inc.) 주식 데이터를 사용합니다. yfinance 패키지를 이용해 데이터를 다운로드하고, 종가(Close) 값만을 사용하여 모델을 학습시킵니다.

import yfinance as yf

# AAPL 주식 데이터 불러오기 (2019-2020)
data = yf.download('AAPL', start='2019-01-01', end='2020-12-31', interval='1d')

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

 

3. Dataloader 구성

주식 데이터는 연속된 값이기 때문에 데이터를 시퀀스로 분할하여 학습해야 합니다. 데이터는 훈련, 검증, 테스트 세트로 나누고, MinMaxScaler를 사용해 데이터를 정규화합니다.

# 데이터셋 분할
train_size = int(0.8 * len(data))
val_size = int(0.1 * len(data))
test_size = len(data) - train_size - val_size

train_data = data.iloc[:train_size]
val_data = data.iloc[train_size:train_size + val_size]
test_data = data.iloc[train_size + val_size:]

# 데이터 정규화
scaler = MinMaxScaler(feature_range=(0, 1))
train_data['Close'] = scaler.fit_transform(train_data[['Close']])
val_data['Close'] = scaler.transform(val_data[['Close']])
test_data['Close'] = scaler.transform(test_data[['Close']])

# 커스텀 데이터셋 정의
class StockDataset(Dataset):
    def __init__(self, data, seq_length):
        self.data = data.values.reshape(-1, 1)
        self.seq_length = seq_length

    def __len__(self):
        return len(self.data) - self.seq_length

    def __getitem__(self, idx):
        x = self.data[idx:idx + self.seq_length]
        y = self.data[idx + self.seq_length, 0]
        return torch.tensor(x, dtype=torch.float32), torch.tensor(y, dtype=torch.float32).unsqueeze(-1)

# 데이터셋 및 데이터로더 생성
seq_length = 10
train_dataset = StockDataset(train_data['Close'], seq_length)
val_dataset = StockDataset(val_data['Close'], seq_length)
test_dataset = StockDataset(test_data['Close'], seq_length)

batch_size = 32
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

 

4. LSTM 모델 설계

RNN 계열의 LSTM(Long Short-Term Memory) 모델을 설계합니다. LSTM은 시계열 데이터에 적합한 모델로, 이전 데이터의 중요한 정보를 기억하고 불필요한 정보를 잊어버리는 메커니즘을 갖고 있어요.

class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers, dropout):
        super(LSTMModel, self).__init__()
        self.lstm = nn.LSTM(input_size=input_size,
                            hidden_size=hidden_size,
                            num_layers=num_layers,
                            dropout=dropout,
                            batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        out, _ = self.lstm(x)
        out = out[:, -1, :]  # 마지막 시퀀스 출력 사용
        out = self.fc(out)
        return out

 

5. 하이퍼파라미터 및 모델 학습

하이퍼파라미터를 설정하고, 모델을 학습시킵니다. 손실 함수로는 MSE(Mean Squared Error)를, 최적화 함수로는 Adam을 사용했습니다.

input_size = 1
hidden_size = 256
output_size = 1
num_layers = 2
dropout = 0.2
Epoch = 200

# 모델 초기화
model = LSTMModel(input_size, hidden_size, output_size, num_layers, dropout).to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 학습 및 평가 함수 정의
def train(model, dataloader, criterion, optimizer):
    model.train()
    total_loss = 0
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(dataloader)

def evaluate(model, dataloader, criterion):
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for inputs, targets in dataloader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            total_loss += loss.item()
    return total_loss / len(dataloader)

# 학습 루프
best_valid_loss = float('inf')
best_epoch = 0

with tqdm(total=Epoch, desc="Training Loop", unit="epoch") as pbar:
    for epoch in range(1, Epoch + 1):
        train_loss = train(model, train_loader, criterion, optimizer)
        valid_loss = evaluate(model, val_loader, criterion)
        pbar.set_postfix({'Train Loss': f'{train_loss:.4f}', 'Validation Loss': f'{valid_loss:.4f}'})
        pbar.update(1)

        # 검증 손실이 가장 낮을 때 모델 저장
        if valid_loss < best_valid_loss:
            best_valid_loss = valid_loss
            best_epoch = epoch
            torch.save({
                'epoch': best_epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'val_loss': best_valid_loss,
            }, "best_model.pth")

print("Training completed.")
print(f"Best model saved from epoch {best_epoch} with validation loss: {best_valid_loss:.4f}")

 

6. 성능 평가 및 시각화

학습이 완료되면 저장된 모델을 불러와 테스트 데이터를 평가하고, RMSE와 R² 스코어를 계산합니다. 마지막으로 예측된 주식 가격과 실제 주식 가격을 시각화합니다.

# 모델 로드 및 성능 평가
checkpoint = torch.load('best_model.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

test_loss = evaluate(model, test_loader, criterion)
rmse = np.sqrt(mean_squared_error(actuals, predictions))
r2 = r2_score(actuals, predictions)

print(f'Test Loss: {test_loss:.4f}')
print(f'RMSE: {rmse:.4f}')
print(f'R² Score: {r2:.4f}')

# 시각화
dates = test_data.index[-len(actuals):]
plt.figure(figsize=(10, 5))
plt.plot(dates, actuals, label='Actual')
plt.plot(dates, predictions, label='Predicted')
plt.title('Actual vs Predicted Stock Prices')
plt.xticks(rotation=45)
plt.legend()
plt.show()