import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import StandardScaler
import yfinance as yf

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

df = yf.download('005930.KS',
                     start='2023-01-01',
                     end='2023-12-31') # 005930 : 삼성전자 주가

df
# 종가 그래프
fig = df['Close'].plot()
X = df.drop('Close', axis=1) # 정답값열 삭제하여 X에 대입.
y = df[['Close']]
X.shape
from sklearn.preprocessing import StandardScaler

Xss = StandardScaler() # 인풋값을 위한.
yss = StandardScaler() # 정답값을 위한.

X_ss = Xss.fit_transform(X)
y_ss = yss.fit_transform(y)

def seq_data(x, y, sequence_length):

  x_seq = []
  y_seq = []
  for i in range(len(x) - sequence_length):
    x_seq.append(x[i: i+sequence_length])
    y_seq.append(y[i+sequence_length])
  # float형 tensor로 변형, gpu사용가능하게 .to(device)를 사용.         # 정답을 열벡터로 줘야해서.
  return torch.FloatTensor(x_seq).to(device), torch.FloatTensor(y_seq).view([-1, 1]).to(device) 

# 200개 데이터를 학습, 나머지를 valid로 사용할 것.
split = 200

# 10일치로 다음날(11일차) 예측
sequence_length = 10

x_seq, y_seq = seq_data(X_ss, y_ss, sequence_length)

x_train_seq = x_seq[:split]
y_train_seq = y_seq[:split]
x_test_seq = x_seq[split:]
y_test_seq = y_seq[split:]
print(x_train_seq.size(), y_train_seq.size())
print(x_test_seq.size(), y_test_seq.size())
train = TensorDataset(x_train_seq, y_train_seq)
test = TensorDataset(x_test_seq, y_test_seq)

batch_size = 16
train_loader = DataLoader(dataset=train, batch_size=batch_size, shuffle=False, drop_last= True)
test_loader = DataLoader(dataset=test, batch_size=batch_size, shuffle=False, drop_last= True)
# GRU 모델 정의
class GRUModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(GRUModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        self.gru = nn.GRU(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)

        out, _ = self.gru(x, h0)
        out = self.fc(out[:, -1, :])
        return out

input_size = x_train_seq.shape[2]
hidden_size = 16
num_layers = 2
output_size = 1 # 1일치 예측

model = GRUModel(input_size, hidden_size, num_layers, output_size).to(device)

# 손실 함수 및 옵티마이저 설정
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 모델 훈련
num_epochs = 100
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, targets in train_loader:
        inputs, targets = inputs.to(device), targets.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')

# 예측 수행 및 시각화
model.eval()
with torch.no_grad():
    train_pred = []
    test_pred = []

    for data in train_loader:
        seq, target = data
        out = model(seq.to(device))
        train_pred += out.cpu().numpy().tolist()

    for data in test_loader:
        seq, target = data
        out = model(seq.to(device))
        test_pred += out.cpu().numpy().tolist()

    total = train_pred + test_pred
    total_inverse = yss.inverse_transform(np.array(total).reshape(-1, 1))

    # 전체 실제 주가 시각화
    plt.figure(figsize=(20,10))
    
    # 학습 데이터 경계선 긋기
    plt.plot(np.ones(100)*len(train_pred), np.linspace(0,1,100), '--', linewidth=0.6)
    
    # 실제 주가
    plt.plot(df['Close'][sequence_length:].values, '--', label='Actual Prices')
    
    # 예측 주가
    plt.plot(np.arange(len(df) - len(total_inverse), len(df)), total_inverse, 'b', linewidth=0.6, label='Predicted Prices')
    
    # y축 범위를 종가 최소 최대에 맞춘다.
    plt.ylim(min(df['Close'].values), max(df['Close'].values))
    plt.title('Actual vs Predicted Samsung Electronics Stock Prices')
    plt.legend()
    plt.xlabel('Time')
    plt.ylabel('Price')
    plt.show()

+ Recent posts