From 7004b281a4c98d67e36b282de713f0974270f040 Mon Sep 17 00:00:00 2001 From: crypto_predictor Date: Thu, 12 Feb 2026 12:43:58 +0800 Subject: [PATCH] refactor: Migrate inference service to use ONNX Runtime --- models/inference.py | 141 +++++++++++++++++++++++--------------------- 1 file changed, 75 insertions(+), 66 deletions(-) diff --git a/models/inference.py b/models/inference.py index d73f283..2c97b1b 100644 --- a/models/inference.py +++ b/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) - model = DLinear(checkpoint['config']) - 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: - self.preprocessors['dlinear'] = pickle.load(f) - except AttributeError: - f.seek(0) - self.preprocessors['dlinear'] = CustomUnpickler(f).load() - print(" ✓ DLinear 已加载") + try: + self.models['dlinear'] = ort.InferenceSession(config['model_path']) + + with open(config['preprocessor_path'], 'rb') as f: + try: + self.preprocessors['dlinear'] = pickle.load(f) + except AttributeError: + f.seek(0) + self.preprocessors['dlinear'] = CustomUnpickler(f).load() + print(" ✓ DLinear (ONNX) 已加载") + except Exception as e: + 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) - model = NBeats(checkpoint['config']) - self.models['nbeats'] = model.to(self.device).eval() - with open(config['preprocessor_path'], 'rb') as f: - try: - self.preprocessors['nbeats'] = pickle.load(f) - except AttributeError: - f.seek(0) - self.preprocessors['nbeats'] = CustomUnpickler(f).load() - print(" ✓ N-BEATS 已加载") + try: + self.models['nbeats'] = ort.InferenceSession(config['model_path']) + with open(config['preprocessor_path'], 'rb') as f: + try: + self.preprocessors['nbeats'] = pickle.load(f) + except AttributeError: + f.seek(0) + self.preprocessors['nbeats'] = CustomUnpickler(f).load() + print(" ✓ N-BEATS (ONNX) 已加载") + 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( - config['model_name'], - trust_remote_code=True - ) - self.models['ttm'] = model.to(self.device).eval() - print(" ✓ TTM-R2 已加载") + self.models['ttm'] = ort.InferenceSession(config['model_path']) + print(" ✓ TTM-R2 (ONNX) 已加载") 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(): - logits = model(X_last).cpu().numpy() + try: + 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(): - output = model(X_last).cpu().numpy() + try: + 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(): - output = model(context_tensor) - pred_price = output.prediction_outputs.cpu().numpy()[0, 0, 0] + try: + input_name = session.get_inputs()[0].name + output_name = session.get_outputs()[0].name - 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) + 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 + # We want the first prediction step [0, 0, 0] + pred_price = output[0, 0, 0] + + 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]: """