subuhui/crpto_predictor
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: refactor model inference architecture

Browse Source 
- 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 2026-02-11 22:07:32 +08:00
1 parent d34b67f522
commit c0a02868ac
14 files changed
+791 -4284
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

api/prediction_service.py
+75 -242
View File
@@ -1,153 +1,42 @@
"""
预测服务 - DLinear 和 TTM-R2 模型预测
预测服务 - 统一模型预测 (DLinear, N-BEATS, TTM-R2)
"""
import os
import sys
import pickle
import threading
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Tuple
from zoneinfo import ZoneInfo
import torch
import numpy as np
import pandas as pd
import pandas_ta as ta
from sklearn.preprocessing import StandardScaler
import numpy as np
# 添加项目路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from data.fetcher import BinanceDataFetcher
from api.models import PredictionModel
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
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)
feature_cols = ['open', 'high', 'low', 'close', 'volume',
'rsi', 'macd', 'macd_hist', 'macd_signal']
if fit:
scaled_features = self.scaler.fit_transform(df[feature_cols])
self.is_fitted = True
self.feature_columns = feature_cols
else:
scaled_features = self.scaler.transform(df[feature_cols])
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
class CustomUnpickler(pickle.Unpickler):
"""自定义反序列化器，用于处理 __main__ 模块的类"""
def find_class(self, module, name):
if module == "__main__":
if name == "ClassificationPreprocessor":
return ClassificationPreprocessor
return super().find_class(module, name)
from models.inference import get_inference_manager
class PredictionService:
"""预测服务 - 管理 DLinear 和 TTM-R2 模型"""
"""预测服务 - 管理多模型预测"""
# 类级别的锁，确保多个请求不会同时生成预测
_prediction_lock = None
_prediction_lock = threading.Lock()
def __init__(self):
import threading
import platform
# 初始化锁
if PredictionService._prediction_lock is None:
PredictionService._prediction_lock = threading.Lock()
# 设备选择：macOS 用 MPS，Windows/Linux 用 CPU
if platform.system() == 'Darwin': # macOS
self.device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
else: # Windows / Linux
self.device = torch.device('cpu')
print(f"使用设备: {self.device}")
self.et_tz = ZoneInfo('America/New_York')
self.fetcher = BinanceDataFetcher()
# 模型和预处理器
self.dlinear_model = None
self.dlinear_preprocessor = None
self.ttm_model = None
self.inference_manager = get_inference_manager()
# 记录正在处理的日期和月份，避免重复处理
self._processing_dates = set()
self._processing_months = set()
self._load_models()
def _load_models(self):
"""加载模型"""
print("正在加载预测模型...")
# 加载 DLinear
from models.dlinear.model import DLinear
dlinear_path = './models/checkpoints/dlinear_classification_2026/best_model.pt'
preprocessor_path = './models/checkpoints/dlinear_classification_2026/preprocessor.pkl'
checkpoint = torch.load(dlinear_path, map_location=self.device, weights_only=False)
self.dlinear_model = DLinear(checkpoint['config'])
self.dlinear_model.load_state_dict(checkpoint['model_state_dict'])
self.dlinear_model = self.dlinear_model.to(self.device).eval()
with open(preprocessor_path, 'rb') as f:
try:
self.dlinear_preprocessor = pickle.load(f)
except AttributeError:
f.seek(0)
self.dlinear_preprocessor = CustomUnpickler(f).load()
print(f" ✓ DLinear 已加载")
# 加载 TTM-R2
from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction
self.ttm_model = TinyTimeMixerForPrediction.from_pretrained(
"ibm-granite/granite-timeseries-ttm-r2",
trust_remote_code=True
)
self.ttm_model = self.ttm_model.to(self.device).eval()
print(f" ✓ TTM-R2 已加载")
print(f" 设备: {self.device}")
# 预加载模型
print("正在预加载预测模型...")
self.inference_manager.load_all_models()
def _get_btc_data(self, days: int = 30) -> pd.DataFrame:
"""获取BTC历史数据"""
@@ -164,83 +53,22 @@ class PredictionService:
return df
def predict_dlinear(self, df: pd.DataFrame) -> Tuple[int, float]:
"""
使用 DLinear 预测下一小时的涨跌
输入: df 包含历史K线数据 (最后一行是最新的已完成K线)
输出: 预测下一小时相对于 df 最后一行收盘价的涨跌
Returns:
(prediction, probability): 预测结果(0=跌, 1=涨)和概率
"""
# 预处理: 使用全部数据，不去掉最后一行
X, _, _ = self.dlinear_preprocessor.preprocess(df, fit=False)
if len(X) == 0:
"""DLinear 预测代理"""
pred, prob = self.inference_manager.predict_dlinear(df)
if pred is None:
return 0, 0.5
# 取最后一个窗口进行预测
X_last = torch.FloatTensor(X[-1:]).to(self.device)
with torch.no_grad():
logits = self.dlinear_model(X_last).cpu().numpy()
prob = 1 / (1 + np.exp(-logits))
prediction = int(prob.item() > 0.5)
probability = float(prob.item())
return prediction, probability
return pred, prob if prob is not None else 0.5
def predict_ttm(self, df: pd.DataFrame) -> Tuple[int, float]:
"""
使用 TTM-R2 预测下一小时的涨跌
输入: df 包含历史K线数据 (最后一行是最新的已完成K线)
输出: 预测下一小时相对于 df 最后一行收盘价的涨跌
Returns:
(prediction, probability): 预测结果(0=跌, 1=涨)和概率
"""
context_length = 512
all_prices = df['close'].values
# 使用最后 context_length 个价格
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)
# 基准价格: 最后一行的收盘价 (即预测目标小时的前一小时收盘价)
base_price = float(df['close'].iloc[-1])
context_tensor = torch.FloatTensor(context_data).to(self.device)
with torch.no_grad():
output = self.ttm_model(context_tensor)
pred_price = output.prediction_outputs.cpu().numpy()[0, 0, 0]
# 判断涨跌: 预测价格 vs 基准价格
prediction = 1 if pred_price > base_price else 0
# 计算伪概率: 基于价格变化百分比
price_pct = (pred_price - base_price) / base_price
# 将价格变化百分比映射到概率 (sigmoid)
probability = 1 / (1 + np.exp(-price_pct * 100))
return prediction, float(probability)
"""TTM-R2 预测代理"""
pred, prob = self.inference_manager.predict_ttm(df)
if pred is None:
return 0, 0.5
return pred, prob if prob is not None else 0.5
def predict_current_hour(self) -> Dict:
"""
预测当前小时的涨跌
逻辑:
- 获取最新的K线数据 (最后一行是最新已完成的K线)
- 预测下一小时相对于最后一行收盘价的涨跌
Returns:
预测结果字典
预测当前小时的涨跌 (多模型集成)
"""
df = self._get_btc_data()
@@ -254,7 +82,6 @@ class PredictionService:
last_close = float(df['close'].iloc[-1])
# 当前正在形成的K线信息
current_open = float(df['open'].iloc[-1])
current_price = float(df['close'].iloc[-1]) # 当前价格
results = {
@@ -266,28 +93,45 @@ class PredictionService:
"models": {}
}
# DLinear 预测 (预测下一小时相对于 last_close 的涨跌)
dlinear_pred, dlinear_prob = self.predict_dlinear(df)
results["models"]["DLinear"] = {
"prediction": dlinear_pred,
"probability": round(dlinear_prob, 4),
"direction": "UP" if dlinear_pred == 1 else "DOWN"
}
# 使用 InferenceManager 获取所有预测
predictions = self.inference_manager.predict_all(df)
# TTM-R2 预测
ttm_pred, ttm_prob = self.predict_ttm(df)
results["models"]["TTM-R2"] = {
"prediction": ttm_pred,
"probability": round(ttm_prob, 4),
"direction": "UP" if ttm_pred == 1 else "DOWN"
}
# 填充结果
if 'dlinear' in predictions:
results["models"]["DLinear"] = {
"prediction": predictions['dlinear']['prediction'],
"probability": round(predictions['dlinear']['probability'], 4),
"direction": "UP" if predictions['dlinear']['prediction'] == 1 else "DOWN"
}
if 'nbeats' in predictions:
results["models"]["N-BEATS"] = {
"prediction": predictions['nbeats']['prediction'],
"probability": round(predictions['nbeats']['probability'], 4),
"direction": "UP" if predictions['nbeats']['prediction'] == 1 else "DOWN"
}
if 'ttm' in predictions:
results["models"]["TTM-R2"] = {
"prediction": predictions['ttm']['prediction'],
"probability": round(predictions['ttm']['probability'], 4),
"direction": "UP" if predictions['ttm']['prediction'] == 1 else "DOWN"
}
if 'ensemble' in predictions and predictions['ensemble']:
ens_pred = predictions['ensemble']['prediction']
results["models"]["Ensemble"] = {
"prediction": ens_pred,
"probability": 0.0, # 集成暂无概率
"direction": "UP" if ens_pred == 1 else "DOWN",
"votes": predictions['ensemble']['votes']
}
return results
def ensure_predictions_for_date(self, date: str) -> None:
"""
确保指定日期有预测数据，如果没有则生成
Args:
date: 日期 (YYYY-MM-DD)
"""
@@ -301,9 +145,15 @@ class PredictionService:
self._processing_dates.add(date)
try:
# 检查是否有缺失的小时
for model_name in ["DLinear", "TTM-R2"]:
missing_hours = PredictionModel.get_missing_hours(date, model_name)
# 检查是否有缺失的小时 (现在只检查 DLinear 和 TTM-R2 作为代表，或者扩展到 N-BEATS)
# 为了保持兼容性，还是主要检查这两个，或者可以增加 N-BEATS
models_to_check = ["DLinear", "TTM-R2", "N-BEATS"]
for model_name in models_to_check:
# 数据库中模型名字对应
db_model_name = model_name
missing_hours = PredictionModel.get_missing_hours(date, db_model_name)
if not missing_hours:
continue
@@ -348,18 +198,6 @@ class PredictionService:
hour: int, model_name: str) -> None:
"""
为指定小时生成预测
预测逻辑:
- 预测目标: {date} {hour}:00 的K线收盘价涨跌
- 输入数据: {hour-1}:00 及以前的K线数据 (不包含目标小时)
- 涨跌判断: 预测价格 vs 前一小时收盘价 (即 {hour-1}:00 的 close)
- 正确率验证: 目标小时的 close vs 前一小时的 close
Args:
df: 历史K线数据 (UTC时间戳)
date: 目标日期 (美东时间 YYYY-MM-DD)
hour: 目标小时 (美东时间 0-23)
model_name: 模型名称
"""
try:
# 1. 构建目标时间 (美东时间 -> UTC)
@@ -374,7 +212,6 @@ class PredictionService:
# 3. 找到目标小时的K线 (用于验证)
target_mask = df['timestamp_dt'] == target_time_utc
if not target_mask.any():
# 目标小时数据不存在，跳过
return
target_row = df[target_mask].iloc[0]
@@ -385,7 +222,6 @@ class PredictionService:
prev_time_utc = target_time_utc - pd.Timedelta(hours=1)
prev_mask = df['timestamp_dt'] == prev_time_utc
if not prev_mask.any():
# 前一小时数据不存在，跳过
return
prev_row = df[prev_mask].iloc[0]
@@ -396,16 +232,23 @@ class PredictionService:
df_input = df[input_mask].copy()
if len(df_input) < 200:
# 数据不足，跳过
return
# 6. 调用模型预测
if model_name == "DLinear":
prediction, probability = self.predict_dlinear(df_input)
else: # TTM-R2
prediction, probability = self.predict_ttm(df_input)
prediction = None
probability = 0.5
# 7. 判断实际涨跌 (目标小时收盘价 vs 前一小时收盘价)
if model_name == "DLinear":
prediction, probability = self.inference_manager.predict_dlinear(df_input)
elif model_name == "TTM-R2":
prediction, probability = self.inference_manager.predict_ttm(df_input)
elif model_name == "N-BEATS":
prediction, probability = self.inference_manager.predict_nbeats(df_input)
if prediction is None:
return
# 7. 判断实际涨跌
actual_result = 1 if target_close > prev_close else 0
# 8. 判断预测是否正确
@@ -434,39 +277,29 @@ class PredictionService:
def ensure_predictions_for_month(self, month: str) -> None:
"""
确保指定月份有预测数据
Args:
month: 月份 (YYYY-MM)
"""
# 使用锁避免并发冲突
with PredictionService._prediction_lock:
# 检查是否正在处理该月份
if month in self._processing_months:
return
# 标记正在处理
self._processing_months.add(month)
try:
# 解析月份
year, mon = map(int, month.split('-'))
# 计算该月的日期范围
from calendar import monthrange
_, last_day = monthrange(year, mon)
# 对于过去的日期，确保有预测数据
today = datetime.now(self.et_tz).date()
for day in range(1, last_day + 1):
date = f"{year}-{mon:02d}-{day:02d}"
date_obj = datetime.strptime(date, '%Y-%m-%d').date()
# 只处理过去和今天的日期
if date_obj <= today:
self.ensure_predictions_for_date(date)
finally:
# 移除处理标记
with PredictionService._prediction_lock:
self._processing_months.discard(month)