xgboost

train/train_xgboost_full.py

@@ -0,0 +1,313 @@
+import argparse
+import csv
+import json
+import math
+import os
+import joblib
+import numpy as np
+
+from xgboost import XGBRegressor
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+from sklearn.metrics import mean_absolute_error
+
+# =========================================================
+# Hyper Parameters
+# =========================================================
+
+TEST_SIZE = 0.2
+RANDOM_STATE = 42
+
+N_ESTIMATORS = 500
+MAX_DEPTH = 6
+LEARNING_RATE = 0.05
+SUBSAMPLE = 0.8
+COLSAMPLE_BYTREE = 0.8
+
+CONTINUE_TRAINING = True
+
+# 예측 성공으로 간주할 허용 오차
+ERROR_TOLERANCE = 0.1
+
+# =========================================================
+# 파일명
+# =========================================================
+
+FEATURES_FILENAME = "features.json"
+MEASURE_FILENAME = "measure.csv"
+
+MODEL_FILENAME = "model.pkl"
+SCALER_FILENAME = "scaler.pkl"
+FEATURE_NAMES_FILENAME = "features.txt"
+
+# =========================================================
+# 무시할 key
+# =========================================================
+
+IGNORE_KEYS = {
+    "songno",
+    "difficulty"
+}
+
+
+def safe_float(value):
+    if value is None:
+        return 0.0
+
+    x = float(value)
+
+    if not math.isfinite(x):
+        return 0.0
+
+    return x
+
+
+def train_model(
+    working_dir: str,
+    data_dir: str
+):
+    # =====================================================
+    # path
+    # =====================================================
+
+    features_path = os.path.join(
+        working_dir,
+        FEATURES_FILENAME
+    )
+
+    measure_path = os.path.join(
+        data_dir,
+        MEASURE_FILENAME
+    )
+
+    model_path = os.path.join(
+        working_dir,
+        MODEL_FILENAME
+    )
+
+    scaler_path = os.path.join(
+        working_dir,
+        SCALER_FILENAME
+    )
+
+    feature_names_path = os.path.join(
+        working_dir,
+        FEATURE_NAMES_FILENAME
+    )
+
+    # =====================================================
+    # features.json
+    # =====================================================
+
+    with open(features_path, "r", encoding="utf-8") as f:
+        feature_data = json.load(f)
+
+    if len(feature_data) == 0:
+        raise ValueError("features.json is empty")
+
+    # =====================================================
+    # feature map
+    # =====================================================
+
+    feature_map = {}
+
+    for item in feature_data:
+        key = (
+            str(item["songno"]),
+            str(item["difficulty"])
+        )
+
+        feature_map[key] = item
+
+    # =====================================================
+    # feature names
+    # =====================================================
+
+    feature_names = sorted([
+        k for k in feature_data[0].keys()
+        if k not in IGNORE_KEYS
+    ])
+
+    # =====================================================
+    # measure.csv
+    # =====================================================
+
+    X = []
+    y = []
+
+    with open(measure_path, "r", encoding="utf-8") as f:
+        reader = csv.reader(f)
+        next(reader, None)
+
+        for row in reader:
+            if len(row) < 3:
+                continue
+
+            measure = safe_float(row[0])
+            songno = str(row[1])
+            diff = str(row[2])
+
+            key = (songno, diff)
+
+            if key not in feature_map:
+                print(
+                    f"[WARN] feature not found: "
+                    f"{songno} {diff}"
+                )
+                continue
+
+            feature_item = feature_map[key]
+
+            features = [
+                safe_float(feature_item.get(k, 0))
+                for k in feature_names
+            ]
+
+            X.append(features)
+            y.append(measure)
+
+    if len(X) == 0:
+        raise ValueError("No training data")
+
+    X = np.array(X, dtype=np.float32)
+    y = np.array(y, dtype=np.float32)
+
+    print(f"Dataset Size: {len(X)}")
+    print(f"Feature Count: {len(feature_names)}")
+
+    # =====================================================
+    # split
+    # =====================================================
+
+    X_train, X_valid, y_train, y_valid = train_test_split(
+        X,
+        y,
+        test_size=TEST_SIZE,
+        random_state=RANDOM_STATE
+    )
+
+    # =====================================================
+    # scaler
+    # =====================================================
+
+    if CONTINUE_TRAINING and os.path.exists(scaler_path):
+        print("Loading existing scaler...")
+
+        scaler = joblib.load(scaler_path)
+
+    else:
+        print("Creating new scaler...")
+
+        scaler = StandardScaler()
+        scaler.fit(X_train)
+
+    X_train = scaler.transform(X_train)
+    X_valid = scaler.transform(X_valid)
+
+    # =====================================================
+    # model
+    # =====================================================
+
+    if CONTINUE_TRAINING and os.path.exists(model_path):
+        print("Loading existing model...")
+
+        model = joblib.load(model_path)
+
+        previous_booster = model.get_booster()
+
+        model.fit(
+            X_train,
+            y_train,
+            xgb_model=previous_booster
+        )
+
+    else:
+        print("Creating new model...")
+
+        model = XGBRegressor(
+            n_estimators=N_ESTIMATORS,
+            max_depth=MAX_DEPTH,
+            learning_rate=LEARNING_RATE,
+            subsample=SUBSAMPLE,
+            colsample_bytree=COLSAMPLE_BYTREE,
+            objective="reg:squarederror",
+            random_state=RANDOM_STATE
+        )
+
+        model.fit(X_train, y_train)
+
+    # =====================================================
+    # evaluate
+    # =====================================================
+
+    pred = model.predict(X_valid)
+
+    mae = mean_absolute_error(y_valid, pred)
+
+    correct = np.sum(
+        np.abs(pred - y_valid) <= ERROR_TOLERANCE
+    )
+
+    accuracy = correct / len(y_valid)
+
+    print(f"\nMAE: {mae:.4f}")
+    print(
+        f"Accuracy "
+        f"(±{ERROR_TOLERANCE}): "
+        f"{accuracy:.4f}"
+    )
+
+    # =====================================================
+    # feature importance
+    # =====================================================
+
+    print("\nFeature Importance:")
+
+    importance = model.feature_importances_
+
+    pairs = list(zip(feature_names, importance))
+    pairs.sort(key=lambda x: x[1], reverse=True)
+
+    for name, score in pairs:
+        print(f"{name:25} {score:.6f}")
+
+    # =====================================================
+    # save
+    # =====================================================
+
+    joblib.dump(model, model_path)
+    joblib.dump(scaler, scaler_path)
+
+    with open(feature_names_path, "w", encoding="utf-8") as f:
+        for name in feature_names:
+            f.write(name + "\n")
+
+    print("\nSaved:")
+    print(model_path)
+    print(scaler_path)
+    print(feature_names_path)
+
+
+# =========================================================
+# main
+# =========================================================
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--workingDir",
+        required=True
+    )
+
+    parser.add_argument(
+        "--dataDir",
+        required=True
+    )
+
+    args = parser.parse_args()
+
+    train_model(
+        args.workingDir,
+        args.dataDir
+    )