307 lines
11 KiB
Python
307 lines
11 KiB
Python
"""
|
|
앙상블 모델 학습 스크립트
|
|
|
|
데이터 소스를 선택하여 LightGBM + PyTorch LSTM + Random Forest 앙상블 모델을 학습합니다.
|
|
- DB 모드(기본): stock_daily_prices 테이블의 KRX 히스토리컬 데이터 사용
|
|
- API 모드: KIS API에서 실시간 차트 데이터 수집
|
|
|
|
사용법:
|
|
python train_ensemble.py # DB 기반 전체 학습
|
|
python train_ensemble.py --source db --limit 20 # DB에서 20종목
|
|
python train_ensemble.py --source api --test # KIS API로 10종목 (기존 방식)
|
|
python train_ensemble.py --min-days 500 # 최소 500일 데이터 있는 종목만
|
|
"""
|
|
|
|
import asyncio
|
|
import sys
|
|
import argparse
|
|
import logging
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
# 프로젝트 루트를 Python 경로에 추가
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
|
|
from app.services.feature_engineering import (
|
|
chart_data_to_dataframe,
|
|
db_data_to_dataframe,
|
|
extract_features,
|
|
prepare_training_data,
|
|
)
|
|
from app.services.ml_ensemble import MLEnsemble
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ─── DB 기반 데이터 수집 ───
|
|
|
|
async def collect_training_data_from_db(limit: int = 200, min_days: int = 200) -> tuple:
|
|
"""DB의 stock_daily_prices에서 학습 데이터 수집
|
|
|
|
Args:
|
|
limit: 최대 종목 수
|
|
min_days: 최소 데이터 일수 (이보다 적으면 제외)
|
|
|
|
Returns:
|
|
(X_tabular, X_seq, y)
|
|
"""
|
|
from sqlalchemy import text
|
|
from app.core.database import AsyncSessionLocal
|
|
|
|
all_X_tabular = []
|
|
all_X_seq = []
|
|
all_y = []
|
|
|
|
async with AsyncSessionLocal() as session:
|
|
# 충분한 데이터가 있는 종목 조회
|
|
result = await session.execute(
|
|
text("""
|
|
SELECT stock_code, stock_name, COUNT(*) as cnt
|
|
FROM stock_daily_prices
|
|
GROUP BY stock_code, stock_name
|
|
HAVING COUNT(*) >= :min_days
|
|
ORDER BY cnt DESC
|
|
LIMIT :limit
|
|
"""),
|
|
{"min_days": min_days, "limit": limit},
|
|
)
|
|
stock_list = [(row[0], row[1], row[2]) for row in result.fetchall()]
|
|
|
|
if not stock_list:
|
|
logger.error(f"DB에 {min_days}일 이상 데이터가 있는 종목이 없습니다")
|
|
logger.error("먼저 python collect_krx_data.py 를 실행하세요")
|
|
return np.array([]), np.array([]), np.array([])
|
|
|
|
logger.info(f"DB에서 {len(stock_list)}개 종목 조회 (최소 {min_days}일)")
|
|
|
|
for i, (stock_code, stock_name, day_count) in enumerate(stock_list):
|
|
logger.info(f"[{i+1}/{len(stock_list)}] {stock_code} ({stock_name}): {day_count}일 데이터")
|
|
|
|
try:
|
|
# DB에서 OHLCV 조회
|
|
rows = await session.execute(
|
|
text("""
|
|
SELECT trade_date, open_price, high_price, low_price,
|
|
close_price, volume
|
|
FROM stock_daily_prices
|
|
WHERE stock_code = :code
|
|
ORDER BY trade_date ASC
|
|
"""),
|
|
{"code": stock_code},
|
|
)
|
|
data = [
|
|
{
|
|
"trade_date": r[0],
|
|
"open_price": r[1],
|
|
"high_price": r[2],
|
|
"low_price": r[3],
|
|
"close_price": r[4],
|
|
"volume": r[5],
|
|
}
|
|
for r in rows.fetchall()
|
|
]
|
|
|
|
if len(data) < min_days:
|
|
continue
|
|
|
|
# DataFrame 변환 및 피처 추출
|
|
df = db_data_to_dataframe(data)
|
|
features_df = extract_features(df)
|
|
|
|
# 학습 데이터 준비
|
|
X_tab, X_seq, y, indices = prepare_training_data(
|
|
df, features_df,
|
|
forecast_days=3,
|
|
up_threshold=0.015,
|
|
down_threshold=-0.015,
|
|
)
|
|
|
|
if len(X_tab) == 0:
|
|
logger.warning(f" {stock_code}: 유효 샘플 없음")
|
|
continue
|
|
|
|
all_X_tabular.append(X_tab)
|
|
all_X_seq.append(X_seq)
|
|
all_y.append(y)
|
|
|
|
logger.info(
|
|
f" {stock_code}: {len(X_tab)}개 샘플 "
|
|
f"(UP:{np.sum(y==2)} HOLD:{np.sum(y==1)} DOWN:{np.sum(y==0)})"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f" {stock_code} 처리 실패: {e}")
|
|
|
|
if not all_X_tabular:
|
|
return np.array([]), np.array([]), np.array([])
|
|
|
|
X_tabular = np.vstack(all_X_tabular)
|
|
X_seq = np.vstack(all_X_seq)
|
|
y = np.concatenate(all_y)
|
|
|
|
return X_tabular, X_seq, y
|
|
|
|
|
|
# ─── KIS API 기반 데이터 수집 (기존 방식) ───
|
|
|
|
async def collect_training_data_from_api(stock_list: list, kis_api) -> tuple:
|
|
"""여러 종목의 차트 데이터를 KIS API에서 수집하여 학습 데이터로 변환"""
|
|
all_X_tabular = []
|
|
all_X_seq = []
|
|
all_y = []
|
|
|
|
for i, (stock_code, stock_name) in enumerate(stock_list):
|
|
logger.info(f"[{i+1}/{len(stock_list)}] 데이터 수집: {stock_code} ({stock_name})")
|
|
|
|
try:
|
|
chart_resp = await kis_api.get_daily_chart(stock_code, "D")
|
|
chart_data = chart_resp.get("output", [])
|
|
|
|
if not chart_data or len(chart_data) < 80:
|
|
logger.warning(f" {stock_code}: 데이터 부족 ({len(chart_data) if chart_data else 0}일)")
|
|
continue
|
|
|
|
df = chart_data_to_dataframe(chart_data)
|
|
features_df = extract_features(df)
|
|
|
|
X_tab, X_seq, y, indices = prepare_training_data(
|
|
df, features_df,
|
|
forecast_days=3,
|
|
up_threshold=0.015,
|
|
down_threshold=-0.015,
|
|
)
|
|
|
|
if len(X_tab) == 0:
|
|
logger.warning(f" {stock_code}: 유효 샘플 없음")
|
|
continue
|
|
|
|
all_X_tabular.append(X_tab)
|
|
all_X_seq.append(X_seq)
|
|
all_y.append(y)
|
|
|
|
logger.info(f" {stock_code}: {len(X_tab)}개 샘플 (UP:{np.sum(y==2)} HOLD:{np.sum(y==1)} DOWN:{np.sum(y==0)})")
|
|
|
|
except Exception as e:
|
|
logger.error(f" {stock_code} 수집 실패: {e}")
|
|
|
|
await asyncio.sleep(3)
|
|
|
|
if not all_X_tabular:
|
|
return np.array([]), np.array([]), np.array([])
|
|
|
|
X_tabular = np.vstack(all_X_tabular)
|
|
X_seq = np.vstack(all_X_seq)
|
|
y = np.concatenate(all_y)
|
|
|
|
return X_tabular, X_seq, y
|
|
|
|
|
|
async def get_stock_list_from_api(kis_api, limit: int = 50) -> list:
|
|
"""거래량 상위 종목 조회 (KIS API)"""
|
|
stocks = []
|
|
|
|
try:
|
|
kospi = await kis_api.get_volume_rank("J")
|
|
if kospi and kospi.get("output"):
|
|
for item in kospi["output"]:
|
|
name = item.get("hts_kor_isnm", "")
|
|
skip_keywords = ["ETF", "레버리지", "인버스", "KODEX", "선물", "TIGER", "ARIRANG"]
|
|
if any(k in name for k in skip_keywords):
|
|
continue
|
|
stocks.append((item["mksc_shrn_iscd"], name))
|
|
except Exception as e:
|
|
logger.error(f"KOSPI 종목 조회 실패: {e}")
|
|
|
|
await asyncio.sleep(3)
|
|
|
|
try:
|
|
kosdaq = await kis_api.get_volume_rank("Q")
|
|
if kosdaq and kosdaq.get("output"):
|
|
for item in kosdaq["output"]:
|
|
name = item.get("hts_kor_isnm", "")
|
|
skip_keywords = ["ETF", "레버리지", "인버스", "KODEX", "선물", "TIGER", "ARIRANG"]
|
|
if any(k in name for k in skip_keywords):
|
|
continue
|
|
stocks.append((item["mksc_shrn_iscd"], name))
|
|
except Exception as e:
|
|
logger.error(f"KOSDAQ 종목 조회 실패: {e}")
|
|
|
|
logger.info(f"총 {len(stocks)}개 종목 조회 완료")
|
|
return stocks[:limit]
|
|
|
|
|
|
# ─── 메인 ───
|
|
|
|
async def main():
|
|
parser = argparse.ArgumentParser(description="앙상블 모델 학습")
|
|
parser.add_argument("--source", choices=["db", "api"], default="db",
|
|
help="데이터 소스 (기본: db)")
|
|
parser.add_argument("--limit", type=int, default=200,
|
|
help="학습 종목 수 (기본: 200)")
|
|
parser.add_argument("--test", action="store_true",
|
|
help="테스트 모드 (10종목)")
|
|
parser.add_argument("--min-days", type=int, default=200,
|
|
help="DB 모드: 최소 데이터 일수 (기본: 200)")
|
|
args = parser.parse_args()
|
|
|
|
limit = 10 if args.test else args.limit
|
|
logger.info(f"데이터 소스: {args.source.upper()}")
|
|
logger.info(f"{'테스트 모드' if args.test else '전체 학습'}: {limit}개 종목")
|
|
|
|
start_time = datetime.now()
|
|
|
|
# 데이터 수집
|
|
if args.source == "db":
|
|
logger.info("=" * 60)
|
|
logger.info("DB(stock_daily_prices)에서 학습 데이터 수집 중...")
|
|
X_tabular, X_seq, y = await collect_training_data_from_db(
|
|
limit=limit, min_days=args.min_days
|
|
)
|
|
else:
|
|
# KIS API 모드
|
|
from app.services.korea_investment_api import KoreaInvestmentAPI
|
|
kis_api = KoreaInvestmentAPI()
|
|
|
|
stock_list = await get_stock_list_from_api(kis_api, limit=limit)
|
|
if not stock_list:
|
|
logger.error("종목 리스트를 가져올 수 없습니다")
|
|
return
|
|
|
|
logger.info("=" * 60)
|
|
logger.info("KIS API에서 학습 데이터 수집 중...")
|
|
X_tabular, X_seq, y = await collect_training_data_from_api(stock_list, kis_api)
|
|
|
|
if len(X_tabular) == 0:
|
|
logger.error("학습 데이터가 없습니다")
|
|
return
|
|
|
|
logger.info(f"총 학습 데이터: {len(X_tabular):,}개 샘플")
|
|
logger.info(f" 클래스 분포: UP={np.sum(y==2):,}, HOLD={np.sum(y==1):,}, DOWN={np.sum(y==0):,}")
|
|
|
|
# 앙상블 학습
|
|
logger.info("=" * 60)
|
|
logger.info("앙상블 모델 학습 중...")
|
|
ensemble = MLEnsemble()
|
|
metrics = ensemble.train(X_tabular, X_seq, y)
|
|
|
|
# 결과 출력
|
|
duration = (datetime.now() - start_time).total_seconds()
|
|
logger.info("=" * 60)
|
|
logger.info("학습 완료!")
|
|
logger.info(f" 소요 시간: {duration:.1f}초 ({duration/60:.1f}분)")
|
|
logger.info(f" 학습 결과:")
|
|
for key, value in metrics.items():
|
|
logger.info(f" {key}: {value}")
|
|
logger.info("=" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|