crypto_predictor
233f0c5d10
- start_system.py: 跨平台进程管理 - Windows 使用 CREATE_NEW_PROCESS_GROUP 和 terminate() - Unix 使用 setsid 和 killpg - 自动检测并启用 Windows ANSI 颜色支持 - 设备选择逻辑优化： - macOS: 优先使用 MPS，回退到 CPU - Windows/Linux: 强制使用 CPU（无 CUDA 依赖） - 影响文件： - api/prediction_service.py - auto_trade_polymarket.py - training/config.py
484 lines
18 KiB
Python
484 lines
18 KiB
Python
"""
|
||
预测服务 - DLinear 和 TTM-R2 模型预测
|
||
"""
|
||
import os
|
||
import sys
|
||
import pickle
|
||
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
|
||
|
||
# 添加项目路径
|
||
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)
|
||
|
||
|
||
class PredictionService:
|
||
"""预测服务 - 管理 DLinear 和 TTM-R2 模型"""
|
||
|
||
# 类级别的锁，确保多个请求不会同时生成预测
|
||
_prediction_lock = None
|
||
|
||
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._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}")
|
||
|
||
def _get_btc_data(self, days: int = 30) -> pd.DataFrame:
|
||
"""获取BTC历史数据"""
|
||
end_date = (datetime.now() + timedelta(days=1)).strftime('%Y-%m-%d')
|
||
start_date = (datetime.now() - timedelta(days=days)).strftime('%Y-%m-%d')
|
||
|
||
df = self.fetcher.fetch_ohlcv(
|
||
symbol='BTC',
|
||
timeframe='1h',
|
||
start_date=start_date,
|
||
end_date=end_date,
|
||
use_cache=False
|
||
)
|
||
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:
|
||
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
|
||
|
||
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)
|
||
|
||
def predict_current_hour(self) -> Dict:
|
||
"""
|
||
预测当前小时的涨跌
|
||
|
||
逻辑:
|
||
- 获取最新的K线数据 (最后一行是最新已完成的K线)
|
||
- 预测下一小时相对于最后一行收盘价的涨跌
|
||
|
||
Returns:
|
||
预测结果字典
|
||
"""
|
||
df = self._get_btc_data()
|
||
|
||
# 获取当前时间 (美东)
|
||
now_et = datetime.now(self.et_tz)
|
||
current_date = now_et.strftime('%Y-%m-%d')
|
||
current_hour = now_et.hour
|
||
|
||
# 最后一个已完成的K线信息 (作为预测基准)
|
||
last_timestamp = df['timestamp'].iloc[-1]
|
||
last_close = float(df['close'].iloc[-1])
|
||
|
||
# 当前正在形成的K线信息
|
||
current_open = float(df['open'].iloc[-1])
|
||
current_price = float(df['close'].iloc[-1]) # 当前价格
|
||
|
||
results = {
|
||
"date": current_date,
|
||
"hour": current_hour,
|
||
"timestamp": str(last_timestamp),
|
||
"base_price": last_close, # 预测基准价格
|
||
"current_price": current_price,
|
||
"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"
|
||
}
|
||
|
||
# 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"
|
||
}
|
||
|
||
return results
|
||
|
||
def ensure_predictions_for_date(self, date: str) -> None:
|
||
"""
|
||
确保指定日期有预测数据，如果没有则生成
|
||
|
||
Args:
|
||
date: 日期 (YYYY-MM-DD)
|
||
"""
|
||
# 使用锁避免并发冲突
|
||
with PredictionService._prediction_lock:
|
||
# 检查是否正在处理该日期
|
||
if date in self._processing_dates:
|
||
return
|
||
|
||
# 标记正在处理
|
||
self._processing_dates.add(date)
|
||
|
||
try:
|
||
# 检查是否有缺失的小时
|
||
for model_name in ["DLinear", "TTM-R2"]:
|
||
missing_hours = PredictionModel.get_missing_hours(date, model_name)
|
||
|
||
if not missing_hours:
|
||
continue
|
||
|
||
print(f"为 {date} 生成 {model_name} 预测 (缺失 {len(missing_hours)} 小时)...")
|
||
|
||
# 获取足够的历史数据
|
||
target_date = datetime.strptime(date, '%Y-%m-%d')
|
||
df = self._get_historical_data(target_date)
|
||
|
||
if df is None or len(df) < 200:
|
||
print(f" 警告: 数据不足，无法为 {date} 生成预测")
|
||
continue
|
||
|
||
# 为每个缺失的小时生成预测
|
||
for hour in missing_hours:
|
||
self._generate_prediction_for_hour(df, date, hour, model_name)
|
||
finally:
|
||
# 移除处理标记
|
||
with PredictionService._prediction_lock:
|
||
self._processing_dates.discard(date)
|
||
|
||
def _get_historical_data(self, target_date: datetime) -> Optional[pd.DataFrame]:
|
||
"""获取指定日期的历史数据"""
|
||
try:
|
||
start_date = (target_date - timedelta(days=30)).strftime('%Y-%m-%d')
|
||
end_date = (target_date + timedelta(days=2)).strftime('%Y-%m-%d')
|
||
|
||
df = self.fetcher.fetch_ohlcv(
|
||
symbol='BTC',
|
||
timeframe='1h',
|
||
start_date=start_date,
|
||
end_date=end_date,
|
||
use_cache=True
|
||
)
|
||
return df
|
||
except Exception as e:
|
||
print(f"获取历史数据失败: {e}")
|
||
return None
|
||
|
||
def _generate_prediction_for_hour(self, df: pd.DataFrame, date: str,
|
||
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)
|
||
target_time_et = pd.Timestamp(f"{date} {hour:02d}:00:00", tz='America/New_York')
|
||
target_time_utc = target_time_et.tz_convert('UTC').tz_localize(None)
|
||
|
||
# 2. 确保 timestamp 列是 datetime 类型
|
||
if 'timestamp_dt' not in df.columns:
|
||
df = df.copy()
|
||
df['timestamp_dt'] = pd.to_datetime(df['timestamp'])
|
||
|
||
# 3. 找到目标小时的K线 (用于验证)
|
||
target_mask = df['timestamp_dt'] == target_time_utc
|
||
if not target_mask.any():
|
||
# 目标小时数据不存在，跳过
|
||
return
|
||
|
||
target_row = df[target_mask].iloc[0]
|
||
target_open = float(target_row['open'])
|
||
target_close = float(target_row['close'])
|
||
|
||
# 4. 找到前一小时的K线 (用于预测输入和基准价格)
|
||
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]
|
||
prev_close = float(prev_row['close'])
|
||
|
||
# 5. 获取预测输入数据: 前一小时及以前的所有数据
|
||
input_mask = df['timestamp_dt'] <= prev_time_utc
|
||
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)
|
||
|
||
# 7. 判断实际涨跌 (目标小时收盘价 vs 前一小时收盘价)
|
||
actual_result = 1 if target_close > prev_close else 0
|
||
|
||
# 8. 判断预测是否正确
|
||
is_correct = 1 if prediction == actual_result else 0
|
||
|
||
# 9. 保存预测结果
|
||
prediction_data = {
|
||
"timestamp": target_time_utc.isoformat(),
|
||
"timestamp_et": target_time_et.isoformat(),
|
||
"date": date,
|
||
"hour": hour,
|
||
"model_name": model_name,
|
||
"prediction": prediction,
|
||
"probability": probability,
|
||
"actual_result": actual_result,
|
||
"is_correct": is_correct,
|
||
"open_price": target_open,
|
||
"close_price": target_close
|
||
}
|
||
|
||
PredictionModel.save_prediction(prediction_data)
|
||
|
||
except Exception as e:
|
||
print(f"生成预测失败 ({date} {hour}:00 {model_name}): {e}")
|
||
|
||
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)
|
||
|
||
|
||
# 全局单例
|
||
_prediction_service: Optional[PredictionService] = None
|
||
|
||
|
||
def get_prediction_service() -> PredictionService:
|
||
"""获取预测服务单例"""
|
||
global _prediction_service
|
||
if _prediction_service is None:
|
||
_prediction_service = PredictionService()
|
||
return _prediction_service
|