refactor: Migrate inference service to use ONNX Runtime

models/inference.py

@@ -8,20 +8,13 @@ import torch
 import numpy as np
 import pandas as pd
 import pandas_ta as ta
+import onnxruntime as ort
 from typing import Dict, List, Optional, Tuple, Union, Any
 from sklearn.preprocessing import StandardScaler
 
 # 添加项目根目录到 sys.path
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-from models.dlinear.model import DLinear
-from models.nbeats.model import NBeats
-try:
-    from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction
-except (ImportError, NameError):
-    # 尝试作为可选依赖处理
-    TinyTimeMixerForPrediction = None
-
 
 
 class CustomUnpickler(pickle.Unpickler):
@@ -171,15 +164,15 @@ class InferenceManager:
         # 配置路径
         self.configs = {
             'dlinear': {
-                'model_path': './models/checkpoints/dlinear_classification_2026/best_model.pt',
+                'model_path': './models/checkpoints/dlinear_classification_2026/model.onnx',
                 'preprocessor_path': './models/checkpoints/dlinear_classification_2026/preprocessor.pkl'
             },
             'nbeats': {
-                'model_path': './models/checkpoints/nbeats_btc_1h/best_model.pt',
+                'model_path': './models/checkpoints/nbeats_btc_1h/model.onnx',
                 'preprocessor_path': './models/checkpoints/nbeats_btc_1h/preprocessor.pkl'
             },
             'ttm': {
-                'model_name': 'ibm-granite/granite-timeseries-ttm-r2'
+                'model_path': './models/checkpoints/ttm_r2_onnx/ttm_r2_crypto_channel_mixing.onnx'
             }
         }
         self.initialized = True
@@ -216,18 +209,18 @@ class InferenceManager:
             print(f"DLinear 模型文件不存在: {config['model_path']}")
             return
 
-        checkpoint = torch.load(config['model_path'], map_location=self.device, weights_only=False)
+        try:
-        model = DLinear(checkpoint['config'])
+            self.models['dlinear'] = ort.InferenceSession(config['model_path'])
-        model.load_state_dict(checkpoint['model_state_dict'])
+            
-        self.models['dlinear'] = model.to(self.device).eval()
+            with open(config['preprocessor_path'], 'rb') as f:
-        
+                try:
-        with open(config['preprocessor_path'], 'rb') as f:
+                    self.preprocessors['dlinear'] = pickle.load(f)
-            try:
+                except AttributeError:
-                self.preprocessors['dlinear'] = pickle.load(f)
+                    f.seek(0)
-            except AttributeError:
+                    self.preprocessors['dlinear'] = CustomUnpickler(f).load()
-                f.seek(0)
+            print("  ✓ DLinear (ONNX) 已加载")
-                self.preprocessors['dlinear'] = CustomUnpickler(f).load()
+        except Exception as e:
-        print("  ✓ DLinear 已加载")
+            print(f"DLinear 加载失败: {e}")
 
     def _load_nbeats(self):
         config = self.configs['nbeats']
@@ -235,31 +228,23 @@ class InferenceManager:
             print(f"N-BEATS 模型文件不存在: {config['model_path']}")
             return
 
-        checkpoint = torch.load(config['model_path'], map_location=self.device, weights_only=False)
+        try:
-        model = NBeats(checkpoint['config'])
+            self.models['nbeats'] = ort.InferenceSession(config['model_path'])
-        self.models['nbeats'] = model.to(self.device).eval()
+            with open(config['preprocessor_path'], 'rb') as f:
-        with open(config['preprocessor_path'], 'rb') as f:
+                 try:
-             try:
+                    self.preprocessors['nbeats'] = pickle.load(f)
-                self.preprocessors['nbeats'] = pickle.load(f)
+                 except AttributeError:
-             except AttributeError:
+                    f.seek(0)
-                f.seek(0)
+                    self.preprocessors['nbeats'] = CustomUnpickler(f).load()
-                self.preprocessors['nbeats'] = CustomUnpickler(f).load()
+            print("  ✓ N-BEATS (ONNX) 已加载")
-        print("  ✓ N-BEATS 已加载")
+        except Exception as e:
+            print(f"N-BEATS 加载失败: {e}")
 
     def _load_ttm(self):
-        # 检查 TTM 类是否已导入
-        if TinyTimeMixerForPrediction is None:
-            print("TTM-R2 模型类未定义 (可能是依赖库缺失)")
-            return
-
         config = self.configs['ttm']
         try:
-            model = TinyTimeMixerForPrediction.from_pretrained(
+            self.models['ttm'] = ort.InferenceSession(config['model_path'])
-                config['model_name'],
+            print("  ✓ TTM-R2 (ONNX) 已加载")
-                trust_remote_code=True
-            )
-            self.models['ttm'] = model.to(self.device).eval()
-            print("  ✓ TTM-R2 已加载")
         except Exception as e:
             print(f"TTM 加载失败: {e}")
 
@@ -272,7 +257,7 @@ class InferenceManager:
             if 'dlinear' not in self.models:
                 return None, None
             
-        model = self.models['dlinear']
+        session = self.models['dlinear']
         preprocessor = self.preprocessors['dlinear']
         
         # 预处理
@@ -282,15 +267,22 @@ class InferenceManager:
             return None, None
             
         # 取最后一个窗口
-        X_last = torch.FloatTensor(X[-1:]).to(self.device)
+        # ONNX 需要 float32
+        X_last = X[-1:].astype(np.float32)
         
-        with torch.no_grad():
+        try:
-            logits = model(X_last).cpu().numpy()
+            input_name = session.get_inputs()[0].name
+            output_name = session.get_outputs()[0].name
+            
+            logits = session.run([output_name], {input_name: X_last})[0]
             prob = 1 / (1 + np.exp(-logits))
             prediction = int(prob.item() > 0.5)
             probability = float(prob.item())
             
-        return prediction, probability
+            return prediction, probability
+        except Exception as e:
+            print(f"DLinear 预测出错: {e}")
+            return None, None
 
     def predict_nbeats(self, df: pd.DataFrame) -> Tuple[Optional[int], Optional[float]]:
         """
@@ -301,7 +293,7 @@ class InferenceManager:
             if 'nbeats' not in self.models:
                 return None, None
             
-        model = self.models['nbeats']
+        session = self.models['nbeats']
         preprocessor = self.preprocessors['nbeats']
         
         # 预处理
@@ -311,11 +303,15 @@ class InferenceManager:
             return None, None
             
         # 取最后一个窗口
-        X_last = torch.FloatTensor(X[-1:]).to(self.device)
+        X_last = X[-1:].astype(np.float32)
         current_price_scaled = X[-1, -1] # 窗口中最后一个点的缩放价格
         
-        with torch.no_grad():
+        try:
-            output = model(X_last).cpu().numpy()
+            input_name = session.get_inputs()[0].name
+            output_name = session.get_outputs()[0].name
+            
+            output = session.run([output_name], {input_name: X_last})[0]
+            
             # output shape: (batch, forecast_horizon) or (batch, 1)
             if output.ndim == 1:
                 pred_price_scaled = output[0]
@@ -327,7 +323,10 @@ class InferenceManager:
             # N-BEATS 回归模型没有直接概率
             probability = 0.6 if prediction == 1 else 0.4 # 占位
             
-        return prediction, probability
+            return prediction, probability
+        except Exception as e:
+            print(f"N-BEATS 预测出错: {e}")
+            return None, None
 
     def predict_ttm(self, df: pd.DataFrame) -> Tuple[Optional[int], Optional[float]]:
         """
@@ -338,7 +337,7 @@ class InferenceManager:
             if 'ttm' not in self.models:
                 return None, None
             
-        model = self.models['ttm']
+        session = self.models['ttm']
         context_length = 512
         
         all_prices = df['close'].values
@@ -353,20 +352,30 @@ class InferenceManager:
             padded[-len(all_prices):] = all_prices
             context_data = padded.reshape(1, -1, 1)
             
-        context_tensor = torch.FloatTensor(context_data).to(self.device)
+        # ONNX input
+        context_data = context_data.astype(np.float32)
         base_price = float(df['close'].iloc[-1])
         
-        with torch.no_grad():
+        try:
-            output = model(context_tensor)
+            input_name = session.get_inputs()[0].name
-            pred_price = output.prediction_outputs.cpu().numpy()[0, 0, 0]
+            output_name = session.get_outputs()[0].name
             
-        prediction = 1 if pred_price > base_price else 0
+            output = session.run([output_name], {input_name: context_data})[0]
-        
+            # output shape might be (batch, prediction_length, features)
-        # 简单的概率模拟
+            # prediction_length=96 usually for TTM zero-shot
-        price_pct = (pred_price - base_price) / base_price
+            # We want the first prediction step [0, 0, 0]
-        probability = 1 / (1 + np.exp(-price_pct * 100))
+            pred_price = output[0, 0, 0]
-        
+            
-        return prediction, float(probability)
+            prediction = 1 if pred_price > base_price else 0
+            
+            # 简单的概率模拟
+            price_pct = (pred_price - base_price) / base_price
+            probability = 1 / (1 + np.exp(-price_pct * 100))
+            
+            return prediction, float(probability)
+        except Exception as e:
+            print(f"TTM 预测出错: {e}")
+            return None, None
 
     def predict_all(self, df: pd.DataFrame) -> Dict[str, Any]:
         """