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

refactor: Migrate inference service to use ONNX Runtime

Browse Source
This commit is contained in:
crypto_predictor committed 2026-02-12 12:43:58 +08:00
1 parent e773775acd
commit 7004b281a4
1 file changed
+75 -66
Download Patch File Download Diff File Expand all files Collapse all files
models/inference.py
+75 -66
View File
@@ -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]:
"""