412 lines
14 KiB
Python
412 lines
14 KiB
Python
import os
|
|
import time
|
|
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
from sklearn.preprocessing import MinMaxScaler
|
|
|
|
from modules.config import Config
|
|
|
|
# cuDNN 벤치마크 활성화 (고정 입력 크기에 대해 최적 커널 자동 선택)
|
|
torch.backends.cudnn.benchmark = True
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(self, hidden_size):
|
|
super(Attention, self).__init__()
|
|
self.attn = nn.Linear(hidden_size, 1)
|
|
|
|
def forward(self, lstm_output):
|
|
attn_weights = torch.softmax(self.attn(lstm_output), dim=1)
|
|
context = torch.sum(attn_weights * lstm_output, dim=1)
|
|
return context, attn_weights
|
|
|
|
|
|
class AdvancedLSTM(nn.Module):
|
|
def __init__(self, input_size=1, hidden_size=512, num_layers=4, output_size=1, dropout=0.3):
|
|
super(AdvancedLSTM, self).__init__()
|
|
self.hidden_size = hidden_size
|
|
self.num_layers = num_layers
|
|
|
|
self.lstm = nn.LSTM(input_size, hidden_size, num_layers,
|
|
batch_first=True, dropout=dropout)
|
|
self.attention = Attention(hidden_size)
|
|
self.fc = nn.Sequential(
|
|
nn.Linear(hidden_size, hidden_size // 2),
|
|
nn.ReLU(),
|
|
nn.Dropout(dropout),
|
|
nn.Linear(hidden_size // 2, hidden_size // 4),
|
|
nn.ReLU(),
|
|
nn.Linear(hidden_size // 4, output_size)
|
|
)
|
|
|
|
def forward(self, x):
|
|
h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
|
|
c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
|
|
lstm_out, _ = self.lstm(x, (h0, c0))
|
|
context, _ = self.attention(lstm_out)
|
|
out = self.fc(context)
|
|
return out
|
|
|
|
|
|
def _unload_ollama():
|
|
"""LSTM 학습 전 Ollama 모델 언로드하여 GPU 메모리 확보"""
|
|
try:
|
|
import requests
|
|
url = f"{Config.OLLAMA_API_URL}/api/generate"
|
|
requests.post(url, json={
|
|
"model": Config.OLLAMA_MODEL,
|
|
"keep_alive": 0
|
|
}, timeout=5)
|
|
print("[AI] Ollama model unloaded (GPU memory freed)")
|
|
time.sleep(1) # 메모리 해제 대기
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _preload_ollama():
|
|
"""LSTM 학습 후 Ollama 모델 다시 로드"""
|
|
try:
|
|
import requests
|
|
url = f"{Config.OLLAMA_API_URL}/api/generate"
|
|
requests.post(url, json={
|
|
"model": Config.OLLAMA_MODEL,
|
|
"prompt": "",
|
|
"keep_alive": "10m"
|
|
}, timeout=10)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _log_gpu_memory(tag=""):
|
|
"""GPU 메모리 사용량 로깅"""
|
|
if torch.cuda.is_available():
|
|
allocated = torch.cuda.memory_allocated(0) / 1024**3
|
|
reserved = torch.cuda.memory_reserved(0) / 1024**3
|
|
print(f"[AI GPU {tag}] Allocated: {allocated:.2f}GB / Reserved: {reserved:.2f}GB")
|
|
|
|
|
|
class PricePredictor:
|
|
"""
|
|
주가 예측 Deep Learning 모델 (GPU 최적화)
|
|
- 전체 학습 데이터를 GPU에 상주 (CPU↔GPU 전송 최소화)
|
|
- Ollama 모델 언로드/리로드로 GPU 메모리 확보
|
|
- Early Stopping + Mixed Precision (FP16)
|
|
- 종목별 모델 체크포인트
|
|
"""
|
|
def __init__(self):
|
|
self.scaler = MinMaxScaler(feature_range=(0, 1))
|
|
|
|
self.hidden_size = 512
|
|
self.num_layers = 4
|
|
|
|
self.model = AdvancedLSTM(input_size=1, hidden_size=self.hidden_size,
|
|
num_layers=self.num_layers, dropout=0.3)
|
|
self.criterion = nn.MSELoss()
|
|
|
|
# CUDA 설정
|
|
self.device = torch.device('cpu')
|
|
self.use_amp = False
|
|
|
|
if torch.cuda.is_available():
|
|
try:
|
|
gpu_name = torch.cuda.get_device_name(0)
|
|
vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
|
|
|
|
self.device = torch.device('cuda')
|
|
self.model.to(self.device)
|
|
|
|
# Mixed Precision (Compute Capability >= 7.0: Volta 이상)
|
|
if torch.cuda.get_device_capability(0)[0] >= 7:
|
|
self.use_amp = True
|
|
|
|
# Warm-up: CUDA 커널 컴파일 유도
|
|
dummy = torch.zeros(1, 60, 1, device=self.device)
|
|
with torch.no_grad():
|
|
_ = self.model(dummy)
|
|
torch.cuda.synchronize()
|
|
|
|
print(f"[AI] GPU Mode: {gpu_name} ({vram_gb:.1f}GB)"
|
|
f" | FP16={'ON' if self.use_amp else 'OFF'}"
|
|
f" | cuDNN Benchmark=ON")
|
|
_log_gpu_memory("init")
|
|
|
|
except Exception as e:
|
|
print(f"[AI] GPU Init Failed ({e}), falling back to CPU")
|
|
self.device = torch.device('cpu')
|
|
self.model.to(self.device)
|
|
else:
|
|
print("[AI] No CUDA GPU detected. Running on CPU.")
|
|
|
|
self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=0.0005, weight_decay=1e-4)
|
|
self.scaler_amp = torch.amp.GradScaler('cuda') if self.use_amp else None
|
|
|
|
self.batch_size = 64
|
|
self.max_epochs = 200
|
|
self.seq_length = 60
|
|
self.patience = 15
|
|
|
|
self.training_status = {
|
|
"is_training": False,
|
|
"loss": 0.0,
|
|
"current_ticker": None
|
|
}
|
|
|
|
@staticmethod
|
|
def verify_hardware():
|
|
if torch.cuda.is_available():
|
|
try:
|
|
gpu_name = torch.cuda.get_device_name(0)
|
|
vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
|
|
print(f"[AI Check] {gpu_name} ({vram_gb:.1f}GB VRAM) | cuDNN={torch.backends.cudnn.is_available()}")
|
|
return True
|
|
except Exception as e:
|
|
print(f"[AI Check] GPU Error: {e}")
|
|
return False
|
|
print("[AI Check] No GPU. CPU Mode.")
|
|
return False
|
|
|
|
def _get_checkpoint_path(self, ticker):
|
|
return os.path.join(Config.MODEL_DIR, f"{ticker}_lstm.pt")
|
|
|
|
def _load_checkpoint(self, ticker):
|
|
path = self._get_checkpoint_path(ticker)
|
|
if os.path.exists(path):
|
|
try:
|
|
checkpoint = torch.load(path, map_location=self.device, weights_only=True)
|
|
self.model.load_state_dict(checkpoint['model_state_dict'])
|
|
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
|
print(f"[AI] Checkpoint loaded: {ticker}")
|
|
return True
|
|
except Exception as e:
|
|
print(f"[AI] Checkpoint load failed ({ticker}): {e}")
|
|
return False
|
|
|
|
def _save_checkpoint(self, ticker, epoch, loss):
|
|
path = self._get_checkpoint_path(ticker)
|
|
try:
|
|
torch.save({
|
|
'model_state_dict': self.model.state_dict(),
|
|
'optimizer_state_dict': self.optimizer.state_dict(),
|
|
'epoch': epoch,
|
|
'loss': loss
|
|
}, path)
|
|
except Exception as e:
|
|
print(f"[AI] Checkpoint save failed ({ticker}): {e}")
|
|
|
|
def train_and_predict(self, prices, forecast_days=1, ticker=None):
|
|
if len(prices) < (self.seq_length + 10):
|
|
return None
|
|
|
|
is_gpu = self.device.type == 'cuda'
|
|
|
|
# --- Ollama 모델 언로드 (GPU 메모리 확보) ---
|
|
if is_gpu:
|
|
_unload_ollama()
|
|
torch.cuda.empty_cache()
|
|
_log_gpu_memory("pre-train")
|
|
|
|
t_start = time.time()
|
|
|
|
# 1. 데이터 전처리 (CPU에서 numpy 작업)
|
|
data = np.array(prices).reshape(-1, 1)
|
|
scaled_data = self.scaler.fit_transform(data)
|
|
|
|
x_seqs, y_seqs = [], []
|
|
for i in range(len(scaled_data) - self.seq_length):
|
|
x_seqs.append(scaled_data[i:i + self.seq_length])
|
|
y_seqs.append(scaled_data[i + self.seq_length])
|
|
|
|
# 2. 텐서 생성 → 즉시 GPU로 이동 (이후 CPU↔GPU 전송 없음)
|
|
x_all = torch.FloatTensor(np.array(x_seqs)).to(self.device)
|
|
y_all = torch.FloatTensor(np.array(y_seqs)).to(self.device)
|
|
|
|
# Validation split (80/20)
|
|
split_idx = int(len(x_all) * 0.8)
|
|
x_train = x_all[:split_idx]
|
|
y_train = y_all[:split_idx]
|
|
x_val = x_all[split_idx:]
|
|
y_val = y_all[split_idx:]
|
|
|
|
dataset_size = len(x_train)
|
|
|
|
# 3. 체크포인트 로드
|
|
has_checkpoint = False
|
|
if ticker:
|
|
has_checkpoint = self._load_checkpoint(ticker)
|
|
max_epochs = 50 if has_checkpoint else self.max_epochs
|
|
|
|
# 4. 학습 (전체 데이터 GPU 상주, DataLoader 미사용)
|
|
self.model.train()
|
|
self.training_status["is_training"] = True
|
|
if ticker:
|
|
self.training_status["current_ticker"] = ticker
|
|
|
|
best_val_loss = float('inf')
|
|
patience_counter = 0
|
|
final_loss = 0.0
|
|
actual_epochs = 0
|
|
|
|
for epoch in range(max_epochs):
|
|
# --- Training (GPU 내에서 셔플 + 미니배치) ---
|
|
perm = torch.randperm(dataset_size, device=self.device)
|
|
x_shuffled = x_train[perm]
|
|
y_shuffled = y_train[perm]
|
|
|
|
epoch_loss = 0.0
|
|
steps = 0
|
|
|
|
for i in range(0, dataset_size, self.batch_size):
|
|
end = min(i + self.batch_size, dataset_size)
|
|
batch_x = x_shuffled[i:end]
|
|
batch_y = y_shuffled[i:end]
|
|
|
|
self.optimizer.zero_grad(set_to_none=True)
|
|
|
|
if self.use_amp:
|
|
with torch.amp.autocast('cuda'):
|
|
outputs = self.model(batch_x)
|
|
loss = self.criterion(outputs, batch_y)
|
|
self.scaler_amp.scale(loss).backward()
|
|
self.scaler_amp.step(self.optimizer)
|
|
self.scaler_amp.update()
|
|
else:
|
|
outputs = self.model(batch_x)
|
|
loss = self.criterion(outputs, batch_y)
|
|
loss.backward()
|
|
self.optimizer.step()
|
|
|
|
epoch_loss += loss.item()
|
|
steps += 1
|
|
|
|
train_loss = epoch_loss / max(1, steps)
|
|
|
|
# --- Validation (GPU에서 직접 수행) ---
|
|
self.model.eval()
|
|
with torch.no_grad():
|
|
if self.use_amp:
|
|
with torch.amp.autocast('cuda'):
|
|
val_out = self.model(x_val)
|
|
val_loss = self.criterion(val_out, y_val).item()
|
|
else:
|
|
val_out = self.model(x_val)
|
|
val_loss = self.criterion(val_out, y_val).item()
|
|
self.model.train()
|
|
|
|
final_loss = train_loss
|
|
actual_epochs = epoch + 1
|
|
|
|
if val_loss < best_val_loss:
|
|
best_val_loss = val_loss
|
|
patience_counter = 0
|
|
else:
|
|
patience_counter += 1
|
|
if patience_counter >= self.patience:
|
|
break
|
|
|
|
self.training_status["is_training"] = False
|
|
self.training_status["loss"] = final_loss
|
|
|
|
if is_gpu:
|
|
torch.cuda.synchronize()
|
|
|
|
elapsed = time.time() - t_start
|
|
print(f"[AI] {ticker or '?'}: {actual_epochs} epochs in {elapsed:.1f}s"
|
|
f" | loss={final_loss:.6f} val={best_val_loss:.6f}"
|
|
f" | device={self.device}")
|
|
|
|
# 5. 체크포인트 저장
|
|
if ticker:
|
|
self._save_checkpoint(ticker, actual_epochs, final_loss)
|
|
|
|
# 6. 예측
|
|
self.model.eval()
|
|
with torch.no_grad():
|
|
last_seq = torch.FloatTensor(
|
|
scaled_data[-self.seq_length:]
|
|
).unsqueeze(0).to(self.device)
|
|
|
|
if self.use_amp:
|
|
with torch.amp.autocast('cuda'):
|
|
predicted_scaled = self.model(last_seq)
|
|
else:
|
|
predicted_scaled = self.model(last_seq)
|
|
|
|
predicted_price = self.scaler.inverse_transform(
|
|
predicted_scaled.cpu().float().numpy())[0][0]
|
|
|
|
# 7. GPU 메모리 정리 + Ollama 리로드
|
|
if is_gpu:
|
|
# 학습 중간 텐서 해제
|
|
del x_all, y_all, x_train, y_train, x_val, y_val
|
|
torch.cuda.empty_cache()
|
|
_log_gpu_memory("post-train")
|
|
_preload_ollama()
|
|
|
|
current_price = prices[-1]
|
|
trend = "UP" if predicted_price > current_price else "DOWN"
|
|
change_rate = ((predicted_price - current_price) / current_price) * 100
|
|
confidence = 1.0 / (1.0 + (final_loss * 100))
|
|
|
|
return {
|
|
"current": current_price,
|
|
"predicted": float(predicted_price),
|
|
"change_rate": round(change_rate, 2),
|
|
"trend": trend,
|
|
"loss": final_loss,
|
|
"confidence": round(confidence, 2),
|
|
"epochs": actual_epochs,
|
|
"device": str(self.device)
|
|
}
|
|
|
|
def batch_predict(self, prices_dict):
|
|
results = {}
|
|
seqs = []
|
|
metas = []
|
|
|
|
for ticker, prices in prices_dict.items():
|
|
if len(prices) < (self.seq_length + 10):
|
|
results[ticker] = None
|
|
continue
|
|
|
|
data = np.array(prices).reshape(-1, 1)
|
|
scaler = MinMaxScaler(feature_range=(0, 1))
|
|
scaled_data = scaler.fit_transform(data)
|
|
|
|
seq = torch.FloatTensor(scaled_data[-self.seq_length:]).unsqueeze(0)
|
|
seqs.append(seq)
|
|
metas.append((ticker, scaler, prices[-1]))
|
|
|
|
if not seqs:
|
|
return results
|
|
|
|
# 배치로 합쳐서 한번에 GPU 추론
|
|
batch = torch.cat(seqs, dim=0).to(self.device)
|
|
|
|
self.model.eval()
|
|
with torch.no_grad():
|
|
if self.use_amp:
|
|
with torch.amp.autocast('cuda'):
|
|
preds = self.model(batch)
|
|
else:
|
|
preds = self.model(batch)
|
|
|
|
preds_cpu = preds.cpu().float().numpy()
|
|
|
|
for i, (ticker, scaler, current_price) in enumerate(metas):
|
|
predicted_price = scaler.inverse_transform(preds_cpu[i:i+1])[0][0]
|
|
trend = "UP" if predicted_price > current_price else "DOWN"
|
|
change_rate = ((predicted_price - current_price) / current_price) * 100
|
|
|
|
results[ticker] = {
|
|
"current": current_price,
|
|
"predicted": float(predicted_price),
|
|
"change_rate": round(change_rate, 2),
|
|
"trend": trend
|
|
}
|
|
|
|
if self.device.type == 'cuda':
|
|
torch.cuda.empty_cache()
|
|
|
|
return results
|