427 lines
15 KiB
Python
427 lines
15 KiB
Python
"""
|
|
模型推理服务 - 统一管理 DLinear, N-BEATS, TTM-R2 模型加载与预测
|
|
"""
|
|
import os
|
|
import sys
|
|
import pickle
|
|
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__))))
|
|
|
|
|
|
|
|
class CustomUnpickler(pickle.Unpickler):
|
|
"""自定义反序列化器，用于处理 __main__ 模块的类重定向"""
|
|
def find_class(self, module, name):
|
|
# 将 __main__ 中的类重定向到当前模块
|
|
if module == "__main__":
|
|
if name == "ClassificationPreprocessor":
|
|
return ClassificationPreprocessor
|
|
if name == "NBeatsPreprocessor":
|
|
return NBeatsPreprocessor
|
|
return super().find_class(module, name)
|
|
|
|
|
|
class ClassificationPreprocessor:
|
|
"""DLinear分类预处理器"""
|
|
|
|
def __init__(self, lookback: int = 168, forecast_horizon: int = 1):
|
|
self.lookback = lookback
|
|
self.forecast_horizon = forecast_horizon
|
|
self.scaler = StandardScaler()
|
|
self.is_fitted = False
|
|
self.feature_columns = ['open', 'high', 'low', 'close', 'volume',
|
|
'rsi', 'macd', 'macd_hist', 'macd_signal']
|
|
|
|
def add_features(self, df: pd.DataFrame) -> pd.DataFrame:
|
|
df = df.copy()
|
|
df['rsi'] = ta.rsi(df['close'], length=14)
|
|
macd = ta.macd(df['close'])
|
|
if macd is not None and not macd.empty:
|
|
df['macd'] = macd.iloc[:, 0]
|
|
df['macd_hist'] = macd.iloc[:, 1]
|
|
df['macd_signal'] = macd.iloc[:, 2]
|
|
return df
|
|
|
|
def create_labels(self, df: pd.DataFrame) -> pd.DataFrame:
|
|
df = df.copy()
|
|
df['future_price'] = df['close'].shift(-self.forecast_horizon)
|
|
df['future_return'] = (df['future_price'] - df['close']) / df['close']
|
|
df['label'] = (df['future_return'] > 0).astype(int)
|
|
return df[:-self.forecast_horizon]
|
|
|
|
def preprocess(self, df: pd.DataFrame, fit: bool = True):
|
|
df = self.add_features(df)
|
|
df = self.create_labels(df)
|
|
df = df.dropna().reset_index(drop=True)
|
|
|
|
if fit:
|
|
scaled_features = self.scaler.fit_transform(df[self.feature_columns])
|
|
self.is_fitted = True
|
|
else:
|
|
scaled_features = self.scaler.transform(df[self.feature_columns])
|
|
|
|
X, y = [], []
|
|
for i in range(len(scaled_features) - self.lookback):
|
|
X.append(scaled_features[i:i + self.lookback])
|
|
y.append(df['label'].iloc[i + self.lookback])
|
|
|
|
return np.array(X), np.array(y), df
|
|
|
|
def preprocess_for_inference(self, df: pd.DataFrame):
|
|
"""推理专用预处理：整个df一次性处理，不生成标签"""
|
|
df = self.add_features(df)
|
|
df = df.dropna().reset_index(drop=True)
|
|
|
|
# 确保有足够的数据
|
|
if len(df) == 0:
|
|
return np.array([]), df, []
|
|
|
|
scaled_features = self.scaler.transform(df[self.feature_columns])
|
|
|
|
X = []
|
|
valid_indices = []
|
|
|
|
if len(scaled_features) < self.lookback:
|
|
return np.array([]), df, []
|
|
|
|
# 注意：这里生成的窗口数量应该是 len(scaled) - lookback + 1
|
|
# 例如 len=168, lookback=168 -> range(1) -> i=0 -> [0:168] (1个窗口)
|
|
for i in range(len(scaled_features) - self.lookback + 1):
|
|
window = scaled_features[i:i + self.lookback]
|
|
X.append(window)
|
|
valid_indices.append(i + self.lookback - 1)
|
|
|
|
return np.array(X), df, valid_indices
|
|
|
|
|
|
class NBeatsPreprocessor:
|
|
"""N-BEATS预处理器"""
|
|
def __init__(self, lookback: int = 168, forecast_horizon: int = 1):
|
|
self.lookback = lookback
|
|
self.forecast_horizon = forecast_horizon
|
|
self.price_scaler = StandardScaler()
|
|
self.is_fitted = False
|
|
|
|
def preprocess(self, df: pd.DataFrame, fit: bool = True):
|
|
df = df.copy()
|
|
prices = df['close'].values.reshape(-1, 1)
|
|
|
|
if fit:
|
|
scaled_prices = self.price_scaler.fit_transform(prices)
|
|
self.is_fitted = True
|
|
else:
|
|
scaled_prices = self.price_scaler.transform(prices)
|
|
|
|
X = []
|
|
valid_indices = []
|
|
|
|
if len(scaled_prices) < self.lookback:
|
|
return np.array([]), df, []
|
|
|
|
for i in range(len(scaled_prices) - self.lookback + 1):
|
|
X.append(scaled_prices[i:i + self.lookback, 0])
|
|
valid_indices.append(i + self.lookback - 1)
|
|
|
|
return np.array(X), df, valid_indices
|
|
|
|
|
|
class InferenceManager:
|
|
"""统一模型推理管理器"""
|
|
|
|
_instance = None
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
if cls._instance is None:
|
|
cls._instance = super(InferenceManager, cls).__new__(cls)
|
|
cls._instance.initialized = False
|
|
return cls._instance
|
|
|
|
def __init__(self, device_str: str = None):
|
|
if self.initialized:
|
|
return
|
|
|
|
import platform
|
|
if device_str:
|
|
self.device = torch.device(device_str)
|
|
elif platform.system() == 'Darwin' and torch.backends.mps.is_available():
|
|
self.device = torch.device('mps')
|
|
else:
|
|
self.device = torch.device('cpu')
|
|
|
|
print(f"InferenceManager 使用设备: {self.device}")
|
|
|
|
self.models = {}
|
|
self.preprocessors = {}
|
|
|
|
# 配置路径
|
|
self.configs = {
|
|
'dlinear': {
|
|
'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/model.onnx',
|
|
'preprocessor_path': './models/checkpoints/nbeats_btc_1h/preprocessor.pkl'
|
|
},
|
|
'ttm': {
|
|
'model_path': './models/checkpoints/ttm_r2_onnx/ttm_r2_crypto_channel_mixing.onnx'
|
|
}
|
|
}
|
|
self.initialized = True
|
|
|
|
def load_model(self, model_name: str):
|
|
"""加载指定模型"""
|
|
if model_name in self.models:
|
|
return
|
|
|
|
print(f"正在加载 {model_name}...")
|
|
try:
|
|
if model_name == 'dlinear':
|
|
self._load_dlinear()
|
|
elif model_name == 'nbeats':
|
|
self._load_nbeats()
|
|
elif model_name == 'ttm':
|
|
self._load_ttm()
|
|
else:
|
|
print(f"未知模型: {model_name}")
|
|
except Exception as e:
|
|
print(f"加载 {model_name} 失败: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
def load_all_models(self):
|
|
"""加载所有模型"""
|
|
self.load_model('dlinear')
|
|
self.load_model('nbeats')
|
|
self.load_model('ttm')
|
|
|
|
def _load_dlinear(self):
|
|
config = self.configs['dlinear']
|
|
if not os.path.exists(config['model_path']):
|
|
print(f"DLinear 模型文件不存在: {config['model_path']}")
|
|
return
|
|
|
|
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']
|
|
if not os.path.exists(config['model_path']):
|
|
print(f"N-BEATS 模型文件不存在: {config['model_path']}")
|
|
return
|
|
|
|
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):
|
|
config = self.configs['ttm']
|
|
try:
|
|
self.models['ttm'] = ort.InferenceSession(config['model_path'])
|
|
print(" ✓ TTM-R2 (ONNX) 已加载")
|
|
except Exception as e:
|
|
print(f"TTM 加载失败: {e}")
|
|
|
|
def predict_dlinear(self, df: pd.DataFrame) -> Tuple[Optional[int], Optional[float]]:
|
|
"""
|
|
DLinear预测 (输入df包含历史数据，预测下一时刻)
|
|
"""
|
|
if 'dlinear' not in self.models:
|
|
self.load_model('dlinear')
|
|
if 'dlinear' not in self.models:
|
|
return None, None
|
|
|
|
session = self.models['dlinear']
|
|
preprocessor = self.preprocessors['dlinear']
|
|
|
|
# 预处理
|
|
X, _, _ = preprocessor.preprocess_for_inference(df)
|
|
|
|
if len(X) == 0:
|
|
return None, None
|
|
|
|
# 取最后一个窗口
|
|
# ONNX 需要 float32
|
|
X_last = X[-1:].astype(np.float32)
|
|
|
|
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
|
|
except Exception as e:
|
|
print(f"DLinear 预测出错: {e}")
|
|
return None, None
|
|
|
|
def predict_nbeats(self, df: pd.DataFrame) -> Tuple[Optional[int], Optional[float]]:
|
|
"""
|
|
N-BEATS预测
|
|
"""
|
|
if 'nbeats' not in self.models:
|
|
self.load_model('nbeats')
|
|
if 'nbeats' not in self.models:
|
|
return None, None
|
|
|
|
session = self.models['nbeats']
|
|
preprocessor = self.preprocessors['nbeats']
|
|
|
|
# 预处理
|
|
X, _, _ = preprocessor.preprocess(df, fit=False)
|
|
|
|
if len(X) == 0:
|
|
return None, None
|
|
|
|
# 取最后一个窗口
|
|
X_last = X[-1:].astype(np.float32)
|
|
current_price_scaled = X[-1, -1] # 窗口中最后一个点的缩放价格
|
|
|
|
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]
|
|
else:
|
|
pred_price_scaled = output[0, 0]
|
|
|
|
prediction = 1 if pred_price_scaled > current_price_scaled else 0
|
|
|
|
# N-BEATS 回归模型没有直接概率
|
|
probability = 0.6 if prediction == 1 else 0.4 # 占位
|
|
|
|
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]]:
|
|
"""
|
|
TTM-R2预测
|
|
"""
|
|
if 'ttm' not in self.models:
|
|
self.load_model('ttm')
|
|
if 'ttm' not in self.models:
|
|
return None, None
|
|
|
|
session = self.models['ttm']
|
|
context_length = 512
|
|
|
|
all_prices = df['close'].values
|
|
if len(all_prices) == 0:
|
|
return None, None
|
|
|
|
# TTM 上下文准备
|
|
if len(all_prices) >= context_length:
|
|
context_data = all_prices[-context_length:].reshape(1, -1, 1)
|
|
else:
|
|
padded = np.full(context_length, all_prices[-1])
|
|
padded[-len(all_prices):] = all_prices
|
|
context_data = padded.reshape(1, -1, 1)
|
|
|
|
# ONNX input
|
|
context_data = context_data.astype(np.float32)
|
|
base_price = float(df['close'].iloc[-1])
|
|
|
|
try:
|
|
input_name = session.get_inputs()[0].name
|
|
output_name = session.get_outputs()[0].name
|
|
|
|
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]:
|
|
"""
|
|
使用所有加载的模型进行预测
|
|
"""
|
|
results = {}
|
|
|
|
# DLinear
|
|
dl_pred, dl_prob = self.predict_dlinear(df)
|
|
if dl_pred is not None:
|
|
results['dlinear'] = {'prediction': dl_pred, 'probability': dl_prob}
|
|
|
|
# N-BEATS
|
|
nb_pred, nb_prob = self.predict_nbeats(df)
|
|
if nb_pred is not None:
|
|
results['nbeats'] = {'prediction': nb_pred, 'probability': nb_prob}
|
|
|
|
# TTM
|
|
ttm_pred, ttm_prob = self.predict_ttm(df)
|
|
if ttm_pred is not None:
|
|
results['ttm'] = {'prediction': ttm_pred, 'probability': ttm_prob}
|
|
|
|
# 集成 (投票)
|
|
votes = []
|
|
if dl_pred is not None: votes.append(dl_pred)
|
|
if nb_pred is not None: votes.append(nb_pred)
|
|
if ttm_pred is not None: votes.append(ttm_pred)
|
|
|
|
if votes:
|
|
# 简单投票: > 50% 胜出
|
|
vote_sum = sum(votes)
|
|
if vote_sum > len(votes) / 2:
|
|
ensemble_pred = 1
|
|
elif vote_sum < len(votes) / 2:
|
|
ensemble_pred = 0
|
|
else:
|
|
# 平票时，如果有 DLinear，优先 DLinear，否则 0 (保守)
|
|
ensemble_pred = dl_pred if dl_pred is not None else 0
|
|
|
|
results['ensemble'] = {'prediction': ensemble_pred, 'votes': votes}
|
|
else:
|
|
results['ensemble'] = None
|
|
|
|
return results
|
|
|
|
def get_inference_manager() -> InferenceManager:
|
|
"""获取单例"""
|
|
return InferenceManager()
|