수정본

2026-02-22 17:23:20 +09:00
parent b147498c31
commit 32e2263f0a
109 changed files with 13979 additions and 4580 deletions
--- a/backend/train_ensemble.py
+++ b/backend/train_ensemble.py
@@ -0,0 +1,306 @@
+"""
+앙상블 모델 학습 스크립트
+
+데이터 소스를 선택하여 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())