Create inference.py

inference.py

@@ -0,0 +1,180 @@
+import torch
+import numpy as np
+import pandas as pd
+import argparse
+import joblib
+import os
+from tqdm import tqdm
+import json  # 👈 JSON 라이브러리 추가
+
+# ⭐️ 수정 사항 1: 경로 문제를 피하기 위해 명시적으로 import 경로 추가
+import sys
+sys.path.append('.')
+
+from models import TimeXer # 사용하는 모델에 맞게 수정
+from utils.metrics import metric # 성능 평가를 위해 추가
+
+# --- 1. 인자 파싱 (수정 없음) ---
+parser = argparse.ArgumentParser(description='Time Series Prediction')
+# ... (이전과 동일한 모든 argparse 코드) ...
+parser.add_argument('--checkpoint_path', type=str, required=True, help='Path to the model checkpoint file (.pth)')
+parser.add_argument('--scaler_path', type=str, required=True, help='Path to the saved scaler file (.gz)')
+parser.add_argument('--predict_input_file', type=str, default=None, help='[Mode 1] Path to the CSV file for single future prediction')
+parser.add_argument('--evaluate_file', type=str, default=None, help='[Mode 2] Path to the CSV file for rolling evaluation')
+parser.add_argument('--model', type=str, default='TimeXer', help='model name')
+parser.add_argument('--task_name', type=str, default='long_term_forecast', help='task name')
+parser.add_argument('--seq_len', type=int, required=True, help='input sequence length')
+parser.add_argument('--pred_len', type=int, required=True, help='prediction sequence length')
+parser.add_argument('--label_len', type=int, required=True, help='start token length')
+parser.add_argument('--features', type=str, required=True, help='M, S, or MS')
+parser.add_argument('--enc_in', type=int, required=True, help='encoder input size')
+parser.add_argument('--dec_in', type=int, required=True, help='decoder input size')
+parser.add_argument('--c_out', type=int, required=True, help='output size')
+parser.add_argument('--d_model', type=int, required=True, help='dimension of model')
+parser.add_argument('--n_heads', type=int, required=True, help='num of heads')
+parser.add_argument('--e_layers', type=int, required=True, help='num of encoder layers')
+parser.add_argument('--d_layers', type=int, required=True, help='num of decoder layers')
+parser.add_argument('--d_ff', type=int, required=True, help='dimension of fcn')
+parser.add_argument('--factor', type=int, required=True, help='attn factor')
+parser.add_argument('--patch_len', type=int, required=True, help='patch length for TimeXer')
+parser.add_argument('--expand', type=int, required=True)
+parser.add_argument('--d_conv', type=int, required=True)
+parser.add_argument('--dropout', type=float, default=0.1, help='dropout')
+parser.add_argument('--embed', type=str, default='timeF', help='time features encoding')
+parser.add_argument('--activation', type=str, default='gelu', help='activation')
+parser.add_argument('--output_attention', action='store_true', help='whether to output attention in ecoder')
+parser.add_argument('--use_norm', type=int, default=1, help='whether to use normalize')
+parser.add_argument('--freq', type=str, default='t', help='freq for time features encoding')
+args = parser.parse_args()
+
+
+# --- 2. 공통 함수: 모델 및 스케일러 로드 (수정 없음) ---
+def load_model_and_scaler(args):
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    # ⭐️ 수정 사항 2: args에 device 정보 추가 (TimeXer 모델 초기화 시 필요할 수 있음)
+    args.device = device
+    model = TimeXer.Model(args).float().to(device)
+    model.load_state_dict(torch.load(args.checkpoint_path, map_location=device))
+    model.eval()
+    scaler = joblib.load(args.scaler_path)
+    # 진행 상황을 stderr로 출력하여 stdout의 JSON 결과와 분리
+    print(f"Using device: {device}", file=sys.stderr)
+    print("Model and scaler loaded successfully.", file=sys.stderr)
+    return model, scaler, device
+
+# --- 3. 모드 1: 단일 미래 예측 함수 (수정 없음) ---
+def predict_future(args, model, scaler, device):
+    # ... (이전과 동일한 코드) ...
+    # 이 함수는 예측 결과(prediction)만 반환하면 됩니다.
+    df_input = pd.read_csv(args.predict_input_file)
+    if 'date' in df_input.columns:
+        df_input = df_input.drop(columns=['date'])
+    raw_input = df_input.tail(args.seq_len).values
+    
+    input_scaled = scaler.transform(raw_input)
+    batch_x = torch.from_numpy(input_scaled).float().unsqueeze(0).to(device)
+    
+    with torch.no_grad():
+        # TimeXer 모델의 forward 함수에 맞게 인자 전달
+        # 여기서는 batch_x만 필요하다고 가정. 필요 시 batch_x_mark 등 추가
+        outputs = model(batch_x)
+        
+    prediction_scaled = outputs.detach().cpu().numpy()[0]
+    
+    # 스케일 복원 로직
+    if args.features == 'MS' and scaler.n_features_in_ > 1:
+        padding = np.zeros((prediction_scaled.shape[0], scaler.n_features_in_ - args.c_out))
+        # 예측 결과를 마지막 feature 자리에 위치
+        prediction_padded = np.concatenate((padding, prediction_scaled), axis=1)
+        prediction = scaler.inverse_transform(prediction_padded)[:, -args.c_out:]
+    else:
+        prediction = scaler.inverse_transform(prediction_scaled)
+
+    return prediction
+
+
+# --- 4. 모드 2: 전체 기간 롤링 평가 함수 (수정 없음) ---
+def evaluate_performance(args, model, scaler, device):
+    # ... (이전과 동일한 코드) ...
+    # 이 함수는 예측값들과 실제값들을 반환하면 됩니다.
+    df_eval = pd.read_csv(args.evaluate_file)
+    if 'date' in df_eval.columns:
+        df_eval = df_eval.drop(columns=['date'])
+    raw_data = df_eval.values
+    data_scaled = scaler.transform(raw_data)
+
+    preds_unscaled = []
+    trues_unscaled = []
+
+    num_samples = len(data_scaled) - args.seq_len - args.pred_len + 1
+    for i in tqdm(range(num_samples), desc="Evaluating", file=sys.stderr):
+        s_begin = i
+        s_end = s_begin + args.seq_len
+        input_scaled = data_scaled[s_begin:s_end]
+        batch_x = torch.from_numpy(input_scaled).float().unsqueeze(0).to(device)
+
+        true_begin = s_end
+        true_end = true_begin + args.pred_len
+        true_scaled = data_scaled[true_begin:true_end]
+        
+        with torch.no_grad():
+            outputs = model(batch_x)
+        
+        pred_scaled = outputs.detach().cpu().numpy()[0]
+
+        if args.features == 'MS' and scaler.n_features_in_ > 1:
+            padding = np.zeros((pred_scaled.shape[0], scaler.n_features_in_ - args.c_out))
+            pred_padded = np.concatenate((padding, pred_scaled), axis=1)
+            pred_unscaled = scaler.inverse_transform(pred_padded)[:, -args.c_out:]
+        else:
+            pred_unscaled = scaler.inverse_transform(pred_scaled)
+        
+        true_unscaled = scaler.inverse_transform(true_scaled)[:, -args.c_out:]
+
+        preds_unscaled.append(pred_unscaled)
+        trues_unscaled.append(true_unscaled)
+
+    return np.array(preds_unscaled), np.array(trues_unscaled)
+
+# --- 5. 메인 로직 (⭐️⭐️⭐️ 이 부분이 완전히 변경되었습니다 ⭐️⭐️⭐️) ---
+if __name__ == '__main__':
+    
+    final_output = {} # 최종 결과를 담을 딕셔너리
+    
+    try:
+        model, scaler, device = load_model_and_scaler(args)
+
+        if args.predict_input_file:
+            print("--- Running in Single Prediction Mode ---", file=sys.stderr)
+            prediction = predict_future(args, model, scaler, device)
+            final_output = {
+                "status": "success",
+                "mode": "single_prediction",
+                "prediction": prediction.flatten().tolist()
+            }
+
+        elif args.evaluate_file:
+            print("--- Running in Rolling Evaluation Mode ---", file=sys.stderr)
+            eval_preds, eval_trues = evaluate_performance(args, model, scaler, device)
+            
+            # 성능 지표 계산
+            mae, mse, _, _, _ = metric(eval_preds, eval_trues)
+            
+            final_output = {
+                "status": "success",
+                "mode": "rolling_evaluation",
+                "mse": mse,
+                "mae": mae,
+                # 전체 예측을 반환하면 너무 크므로, 샘플만 반환하거나 필요한 정보만 반환
+                "prediction_samples": [p.flatten().tolist() for p in eval_preds[:5]]
+            }
+            
+        else:
+            final_output = {"status": "error", "message": "No mode selected. Use --predict_input_file or --evaluate_file."}
+
+    except Exception as e:
+        final_output = {"status": "error", "message": str(e)}
+
+    # 최종 결과를 JSON 문자열로 표준 출력(stdout)에 프린트합니다.
+    # 이 출력을 app.py가 읽어서 API 응답으로 사용합니다.
+    print(json.dumps(final_output, indent=2))