xgboost

train/train_xgboost_pick.py

@@ -0,0 +1,369 @@
+import argparse
+import csv
+import json
+import math
+import os
+import random
+import joblib
+import numpy as np
+
+from xgboost import XGBRegressor
+from sklearn.preprocessing import StandardScaler
+from sklearn.metrics import mean_absolute_error
+
+# =========================================================
+# Hyper Parameters
+# =========================================================
+
+TRAIN_SIZE = 0
+VALID_SIZE = 0
+
+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.5
+
+# =========================================================
+# 파일명
+# =========================================================
+
+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"
+}
+
+
+# =========================================================
+# safe float
+# =========================================================
+
+def safe_float(value):
+    if value is None:
+        return 0.0
+
+    x = float(value)
+
+    if not math.isfinite(x):
+        return 0.0
+
+    return x
+
+
+# =========================================================
+# train
+# =========================================================
+
+def train_model(
+    working_dir: str,
+    data_dir: str
+):
+    random.seed(RANDOM_STATE)
+
+    # =====================================================
+    # 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
+    ])
+
+    # =====================================================
+    # dataset build
+    # =====================================================
+
+    dataset = []
+
+    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
+            ]
+
+            dataset.append((
+                features,
+                measure
+            ))
+
+    # =====================================================
+    # shuffle
+    # =====================================================
+
+    random.shuffle(dataset)
+
+    required_size = TRAIN_SIZE + VALID_SIZE
+
+    if len(dataset) < required_size:
+        raise ValueError(
+            f"Not enough dataset "
+            f"({len(dataset)} < {required_size})"
+        )
+
+    # =====================================================
+    # split
+    # =====================================================
+
+    train_dataset = dataset[:TRAIN_SIZE]
+    valid_dataset = dataset[
+        TRAIN_SIZE:
+        TRAIN_SIZE + VALID_SIZE
+    ]
+
+    X_train = np.array(
+        [x for x, _ in train_dataset],
+        dtype=np.float32
+    )
+
+    y_train = np.array(
+        [y for _, y in train_dataset],
+        dtype=np.float32
+    )
+
+    X_valid = np.array(
+        [x for x, _ in valid_dataset],
+        dtype=np.float32
+    )
+
+    y_valid = np.array(
+        [y for _, y in valid_dataset],
+        dtype=np.float32
+    )
+
+    print(f"Train Size: {len(X_train)}")
+    print(f"Valid Size: {len(X_valid)}")
+    print(f"Feature Count: {len(feature_names)}")
+
+    # =====================================================
+    # 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
+    )
+    
+    parser.add_argument(
+        "--trainSize",
+        required=True
+    )
+    
+    parser.add_argument(
+        "--validSize",
+        required=True
+    )
+
+    args = parser.parse_args()
+    
+    TRAIN_SIZE = int(args.trainSize)
+    VALID_SIZE = int(args.validSize)
+
+    train_model(
+        args.workingDir,
+        args.dataDir
+    )