feat: refactor model inference architecture
- Unified inference logic into models/inference.py - Added InferenceManager for handling DLinear, N-BEATS, and TTM-R2 models - Updated prediction_service.py and auto_trade_polymarket.py to use the new manager - Fixed dependency issues (torch, transformers) - Resolved logger duplication issue - Verified with test_models_feb_2026.py
This commit is contained in:
crypto_predictor committed
1 parent
d34b67f522
commit
c0a02868ac
14 files changed
+791
-4284
No files matched your search
@@ -0,0 +1,417 @@
|
||||
"""
|
||||
模型推理服务 - 统一管理 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
|
||||
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):
|
||||
"""自定义反序列化器，用于处理 __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/best_model.pt',
|
||||
'preprocessor_path': './models/checkpoints/dlinear_classification_2026/preprocessor.pkl'
|
||||
},
|
||||
'nbeats': {
|
||||
'model_path': './models/checkpoints/nbeats_btc_1h/best_model.pt',
|
||||
'preprocessor_path': './models/checkpoints/nbeats_btc_1h/preprocessor.pkl'
|
||||
},
|
||||
'ttm': {
|
||||
'model_name': 'ibm-granite/granite-timeseries-ttm-r2'
|
||||
}
|
||||
}
|
||||
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
|
||||
|
||||
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 已加载")
|
||||
|
||||
def _load_nbeats(self):
|
||||
config = self.configs['nbeats']
|
||||
if not os.path.exists(config['model_path']):
|
||||
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 已加载")
|
||||
|
||||
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 已加载")
|
||||
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
|
||||
|
||||
model = self.models['dlinear']
|
||||
preprocessor = self.preprocessors['dlinear']
|
||||
|
||||
# 预处理
|
||||
X, _, _ = preprocessor.preprocess_for_inference(df)
|
||||
|
||||
if len(X) == 0:
|
||||
return None, None
|
||||
|
||||
# 取最后一个窗口
|
||||
X_last = torch.FloatTensor(X[-1:]).to(self.device)
|
||||
|
||||
with torch.no_grad():
|
||||
logits = model(X_last).cpu().numpy()
|
||||
prob = 1 / (1 + np.exp(-logits))
|
||||
prediction = int(prob.item() > 0.5)
|
||||
probability = float(prob.item())
|
||||
|
||||
return prediction, probability
|
||||
|
||||
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
|
||||
|
||||
model = self.models['nbeats']
|
||||
preprocessor = self.preprocessors['nbeats']
|
||||
|
||||
# 预处理
|
||||
X, _, _ = preprocessor.preprocess(df, fit=False)
|
||||
|
||||
if len(X) == 0:
|
||||
return None, None
|
||||
|
||||
# 取最后一个窗口
|
||||
X_last = torch.FloatTensor(X[-1:]).to(self.device)
|
||||
current_price_scaled = X[-1, -1] # 窗口中最后一个点的缩放价格
|
||||
|
||||
with torch.no_grad():
|
||||
output = model(X_last).cpu().numpy()
|
||||
# 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
|
||||
|
||||
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
|
||||
|
||||
model = 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)
|
||||
|
||||
context_tensor = torch.FloatTensor(context_data).to(self.device)
|
||||
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]
|
||||
|
||||
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)
|
||||
|
||||
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()
|
||||
Reference in new issue
Block a user