369 lines
8.3 KiB
Python
369 lines
8.3 KiB
Python
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
|
|
) |