.

script/train.py

@@ -1,67 +1,152 @@
 import pandas as pd
-import json, os, sys, joblib
+import json, os, sys
 import numpy as np
-from sklearn.ensemble import GradientBoostingRegressor
-from sklearn.model_selection import train_test_split
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from sklearn.preprocessing import StandardScaler
 
-def train(script_dir, working_dir, data_dir, train_count, val_count, margin):
-    features = pd.read_json(os.path.join(working_dir, 'features.json'))
-    measures = pd.read_csv(os.path.join(data_dir, 'measure.csv'))
-    
-    features['songno'] = features['songno'].astype(int)
-    measures['songno'] = measures['songno'].astype(int)
-    measures['diff'] = measures['diff'].replace('ura', 'edit')
-    
-    df = pd.merge(features, measures, on=['songno', 'diff'])
-    
-    with open(os.path.join(script_dir, 'factor.json'), 'r') as f:
-        weights = json.load(f)
-    for col in ['physical_density', 'stamina_requirement', 'pattern_complexity', 'rhythmic_complexity', 'reading_gimmick']:
-        df[col] = df[col] * weights.get(col, 1.0)
+MAX_MARGIN_LIMIT = 3
 
-    X_cols = ['physical_density', 'stamina_requirement', 'pattern_complexity', 'rhythmic_complexity', 'reading_gimmick']
-    model_path = os.path.join(working_dir, 'model.pkl')
-    model = joblib.load(model_path) if os.path.exists(model_path) else GradientBoostingRegressor(n_estimators=200, learning_rate=0.05, max_depth=3)
+class DifficultyNet(nn.Module):
+    def __init__(self, input_dim):
+        super(DifficultyNet, self).__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, 128),
+            nn.BatchNorm1d(128),
+            nn.ReLU(),
+            nn.Dropout(0.4),
+            nn.Linear(128, 64),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(64, 32),
+            nn.ReLU(),
+            nn.Linear(32, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+def train(script_dir, working_dir, data_dir, train_count, val_count, margin, val_iterations):
+    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+    print(f"🚀 Using Device: {device}")
+
+    # 데이터 로드 및 병합
+    factors_path = os.path.join(working_dir, 'factors.json')
+    if not os.path.exists(factors_path):
+        os.system(f"bun run {os.path.join(script_dir, 'factorize.ts')} {data_dir} {working_dir}")
+
+    df_feat = pd.read_json(factors_path)
+    df_meas = pd.read_csv(os.path.join(data_dir, 'measure.csv'))
+    df_meas['diff'] = df_meas['diff'].replace('ura', 'edit')
+    df_feat['songno'], df_meas['songno'] = df_feat['songno'].astype(int), df_meas['songno'].astype(int)
+    df = pd.merge(df_feat, df_meas, on=['songno', 'diff'])
+
+    exclude_cols = ['songno', 'diff', 'title', 'course', '상수', 'predicted_measure', 'error']
+    X_cols = [c for c in df_feat.columns if c not in exclude_cols]
+    
+    scaler = StandardScaler()
+    df[X_cols] = scaler.fit_transform(df[X_cols])
+    with open(os.path.join(working_dir, 'scaler.json'), 'w') as f:
+        json.dump({'mean': scaler.mean_.tolist(), 'std': scaler.scale_.tolist(), 'cols': X_cols}, f)
+
+    model_path = os.path.join(working_dir, 'model.pth')
+    margin = float(margin)
+    val_iterations = int(val_iterations)
+    
+    model = DifficultyNet(len(X_cols)).to(device)
+    attempt = 1
 
-    iteration = 1
     while True:
-        # 3. 데이터 샘플링
-        df_sample = df.sample(n=int(train_count) + int(val_count))
-        train_df, val_df = train_test_split(df_sample, test_size=int(val_count))
+        print(f"\n[Attempt {attempt}] " + "-"*40)
+        if os.path.exists(model_path):
+            try: model.load_state_dict(torch.load(model_path, map_location=device))
+            except: pass
 
-        X_train, y_train = train_df[X_cols], train_df['상수']
-        X_val, y_val = val_df[X_cols], val_df['상수']
+        train_df = df.sample(n=min(int(train_count), len(df)))
+        X_train = torch.FloatTensor(train_df[X_cols].values).to(device)
+        y_train = torch.FloatTensor(((train_df['상수'].values - 1) / 11)).unsqueeze(1).to(device)
 
-        # 4. 학습 (최대 10회)
-        for attempt in range(1, 11):
-            model.fit(X_train, y_train)
-            train_err = np.max(np.abs(np.clip(model.predict(X_train), 1.0, 12.0) - y_train))
-            print(f"Iteration {iteration} - Attempt {attempt} - Train Error: {train_err:.4f}")
-            if train_err <= float(margin): break
-            model.set_params(n_estimators=model.n_estimators + 50)
+        base_lr = 0.001
+        optimizer = optim.Adam(model.parameters(), lr=base_lr, weight_decay=1e-5)
+        criterion = nn.L1Loss()
 
-        # 5. 검증
-        pred_val = np.clip(model.predict(X_val), 1.0, 12.0)
-        val_errors = np.abs(pred_val - y_val)
-        
-        # 6. 검증 실패 시 재시도
-        if np.any(val_errors > float(margin)):
-            print(f"Validation failed (max error: {np.max(val_errors):.4f}). Retrying...")
-            iteration += 1
-            continue
+        model.train()
+        epoch = 0
+        while True:
+            optimizer.zero_grad()
+            preds_train = model(X_train)
+            loss = criterion(preds_train, y_train) # MSELoss 권장
+            loss.backward()
+            optimizer.step()
             
-        val_result = pd.DataFrame({
-            'songno': val_df['songno'],
-            'difficulty': val_df['diff'],
-            'measure': y_val,
-            'predicted_measure': pred_val,
-            'error': val_errors
-        })
-        val_result.to_csv(os.path.join(working_dir, f'validate_result_{iteration}.csv'), index=False)
-        val_result.to_csv(os.path.join(working_dir, 'validate_result.csv'), index=False)
-        break
+            # 현재 학습 데이터에 대한 오차 분석
+            with torch.no_grad():
+                diff = torch.abs((preds_train * 11 + 1) - (y_train * 11 + 1))
+                train_mae = torch.mean(diff).item()
+                train_max = torch.max(diff).item()
+            
+            if epoch % 1000 == 0:
+                print(f"   - Ep {epoch:5d} | MAE: {train_mae:.4f} | MAX: {train_max:.4f}")
+            
+            # ✅ 조건 강화: MAE뿐만 아니라 MAX도 어느 정도 잡혔을 때만 검증으로 이동
+            if train_mae < margin and train_max < (margin * MAX_MARGIN_LIMIT): 
+                print(f"   ✅ Train Goal Reached (MAE: {train_mae} < {margin:.4f}, MAX: {train_max:.4f} < {(margin * MAX_MARGIN_LIMIT):.4f}). Moving to Val.")
+                break
+            
+            epoch += 1
+            if epoch > 50000:
+                print("   ⚠️ Timed out. Resampling...")
+                break
+        
+        if train_mae >= margin:
+            torch.save(model.state_dict(), model_path)
+            attempt += 1
+            continue
 
-    joblib.dump(model, model_path)
+        # 3. 검증 단계
+        print(f"3. Validating {val_iterations} iterations...")
+        model.eval()
+        all_passed = True
+        with torch.no_grad():
+            for i in range(1, val_iterations + 1):
+                val_df = df.sample(n=min(int(val_count), len(df)))
+                X_val = torch.FloatTensor(val_df[X_cols].values).to(device)
+                y_val_raw = val_df['상수'].values
+                preds = model(X_val) * 11 + 1
+                y_val_tensor = torch.FloatTensor(y_val_raw).unsqueeze(1).to(device)
+                diff_tensor = torch.abs(preds - y_val_tensor)
+                mae = torch.mean(diff_tensor).item()
+                max_error = torch.max(diff_tensor).item()
+                
+
+                # csv 저장
+                # preds를 CPU 넘파이로 변환하고 1차원으로 펴주는 과정이 포함되어야 합니다.
+                preds_np = preds.detach().cpu().numpy().flatten() if torch.is_tensor(preds) else preds
+                
+                res_df = val_df[['songno', 'diff', '상수']].copy()
+                res_df['predicted_measure'] = preds_np
+                res_df['error'] = np.abs(preds_np - y_val_raw)
+                
+                output_file = os.path.join(working_dir, f'validate_result_{i}.csv')
+                res_df.to_csv(output_file, index=False)
+
+                if mae <= margin and max_error <= (margin * MAX_MARGIN_LIMIT):
+                    print(f"   [Iter {i}] ✅ PASS (MAE: {mae:.4f}, MAX: {max_error:.4f})")
+                else:
+                    print(f"   [Iter {i}] ❌ FAIL (MAE: {mae:.4f} > {margin}, MAX: {max_error:.4f} > {margin * MAX_MARGIN_LIMIT})")
+                    all_passed = False
+                    break 
+
+        if all_passed:
+            print(f"\n✨ Final Success!")
+            torch.save(model.state_dict(), model_path)
+            return
+        else:
+            torch.save(model.state_dict(), model_path)
+            attempt += 1
 
 if __name__ == "__main__":
-    train(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5], sys.argv[6])
+    if len(sys.argv) < 8: sys.exit(1)
+    train(*sys.argv[1:8])