crypto_predictor
1bd9d25f60
- 删除测试脚本: test_2026_classification.py, test_lagllama.py, test_timesfm.py, test_ttm.py - 删除2019相关文件: analyze_2019_*, apply_2019_*, dlinear_2019_*, three_models_2019_* - 添加集成模型相关代码和分析脚本 - 添加N-BEATS模型和CatBoost分类模型
296 lines
9.7 KiB
Python
296 lines
9.7 KiB
Python
"""将Google TimesFM 2.5模型导出为ONNX格式 - 使用正确的API"""
|
||
|
||
import sys
|
||
import os
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
import torch
|
||
import numpy as np
|
||
|
||
|
||
def export_timesfm_to_onnx():
|
||
print("="*80)
|
||
print("导出TimesFM 2.5 (200M) 模型为ONNX格式")
|
||
print("="*80)
|
||
|
||
# ========== 步骤1: 加载TimesFM模型 ==========
|
||
print("\n[1/6] 加载TimesFM 2.5模型...")
|
||
|
||
try:
|
||
import timesfm
|
||
|
||
# 使用正确的from_pretrained方法加载模型
|
||
print(" 从HuggingFace加载预训练模型...")
|
||
print(" 模型ID: google/timesfm-2.5-200m-pytorch")
|
||
|
||
model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
|
||
"google/timesfm-2.5-200m-pytorch",
|
||
torch_compile=False # 导出ONNX时不使用torch.compile
|
||
)
|
||
|
||
print("✓ TimesFM 2.5模型已加载")
|
||
|
||
except ImportError as e:
|
||
print(f"✗ TimesFM未安装或导入失败: {e}")
|
||
print(" 请运行: pip install git+https://github.com/google-research/timesfm.git")
|
||
return
|
||
except Exception as e:
|
||
print(f"✗ 加载模型失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return
|
||
|
||
# ========== 步骤2: 编译模型配置 ==========
|
||
print("\n[2/6] 配置模型...")
|
||
|
||
try:
|
||
# 配置模型
|
||
model.compile(
|
||
timesfm.ForecastConfig(
|
||
max_context=512, # BTC历史数据512小时
|
||
max_horizon=24, # 预测未来24小时
|
||
normalize_inputs=True,
|
||
use_continuous_quantile_head=True,
|
||
force_flip_invariance=True,
|
||
infer_is_positive=True,
|
||
fix_quantile_crossing=True,
|
||
)
|
||
)
|
||
|
||
print("✓ 模型配置完成")
|
||
print(f" 最大上下文长度: 512")
|
||
print(f" 最大预测长度: 24")
|
||
print(f" 归一化输入: True")
|
||
print(f" 使用分位数预测: True")
|
||
|
||
except Exception as e:
|
||
print(f"✗ 配置失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return
|
||
|
||
# ========== 步骤3: 测试模型 ==========
|
||
print("\n[3/6] 测试模型推理...")
|
||
|
||
try:
|
||
# 创建测试数据
|
||
test_input = [
|
||
np.linspace(85000, 90000, 512), # 模拟BTC价格从85k到90k
|
||
np.sin(np.linspace(0, 20, 512)) * 1000 + 87000, # 模拟震荡价格
|
||
]
|
||
|
||
# 测试推理
|
||
point_forecast, quantile_forecast = model.forecast(
|
||
horizon=24,
|
||
inputs=test_input,
|
||
)
|
||
|
||
print("✓ 模型推理成功")
|
||
print(f" 点预测形状: {point_forecast.shape}") # (2, 24)
|
||
print(f" 分位数预测形状: {quantile_forecast.shape}") # (2, 24, 11)
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ 模型测试失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
print(" 继续尝试导出...")
|
||
|
||
# ========== 步骤4: 提取核心模型并导出ONNX ==========
|
||
print("\n[4/6] 提取核心模型...")
|
||
|
||
try:
|
||
# TimesFM的核心模型在model.model中
|
||
core_model = model.model
|
||
core_model.eval()
|
||
|
||
print("✓ 核心模型已提取")
|
||
print(f" 模型类型: {type(core_model)}")
|
||
|
||
# 创建ONNX导出包装器
|
||
class TimesFMONNXWrapper(torch.nn.Module):
|
||
"""TimesFM ONNX导出包装器"""
|
||
def __init__(self, timesfm_core_model):
|
||
super().__init__()
|
||
self.model = timesfm_core_model
|
||
|
||
def forward(self, inputs, masks=None):
|
||
"""
|
||
Args:
|
||
inputs: (batch_size, context_length) - 输入时间序列
|
||
masks: (batch_size, context_length) - 可选mask，1=有效，0=无效
|
||
Returns:
|
||
point_forecast: (batch_size, horizon_length) - 点预测
|
||
"""
|
||
# 如果没有提供masks，创建全1的mask
|
||
if masks is None:
|
||
masks = torch.ones_like(inputs)
|
||
|
||
# 调用TimesFM核心模型
|
||
# 注意：这里需要根据实际的forward签名调整
|
||
outputs, _ = self.model(inputs, masks, decode_caches=None)
|
||
|
||
# outputs包含: (input_embeddings, output_embeddings, output_ts, output_quantile_spread)
|
||
# output_ts是点预测
|
||
point_forecast = outputs[2]
|
||
|
||
return point_forecast
|
||
|
||
wrapped_model = TimesFMONNXWrapper(core_model)
|
||
wrapped_model.eval()
|
||
|
||
print("✓ ONNX包装器已创建")
|
||
|
||
except Exception as e:
|
||
print(f"✗ 提取核心模型失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return
|
||
|
||
# ========== 步骤5: 导出为ONNX ==========
|
||
print("\n[5/6] 导出模型为ONNX格式...")
|
||
|
||
onnx_path = "./models/timesfm_2p5_200m.onnx"
|
||
os.makedirs("./models", exist_ok=True)
|
||
|
||
try:
|
||
# 准备示例输入
|
||
batch_size = 1
|
||
context_len = 512
|
||
|
||
dummy_input = torch.randn(batch_size, context_len, dtype=torch.float32)
|
||
dummy_mask = torch.ones(batch_size, context_len, dtype=torch.float32)
|
||
|
||
print(f" 示例输入形状: inputs={dummy_input.shape}, masks={dummy_mask.shape}")
|
||
print(" 开始导出ONNX...")
|
||
|
||
# 导出为ONNX
|
||
torch.onnx.export(
|
||
wrapped_model,
|
||
(dummy_input, dummy_mask),
|
||
onnx_path,
|
||
export_params=True,
|
||
opset_version=14,
|
||
do_constant_folding=True,
|
||
input_names=['inputs', 'masks'],
|
||
output_names=['point_forecast'],
|
||
dynamic_axes={
|
||
'inputs': {0: 'batch_size', 1: 'context_length'},
|
||
'masks': {0: 'batch_size', 1: 'context_length'},
|
||
'point_forecast': {0: 'batch_size', 1: 'horizon_length'}
|
||
},
|
||
verbose=False
|
||
)
|
||
|
||
print(f"✓ ONNX模型已导出")
|
||
print(f" 保存路径: {onnx_path}")
|
||
print(f" 文件大小: {os.path.getsize(onnx_path) / (1024**2):.2f} MB")
|
||
|
||
except Exception as e:
|
||
print(f"✗ ONNX导出失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
# 尝试简化版本
|
||
print("\n尝试简化导出方案...")
|
||
try:
|
||
# 只导出单输入版本
|
||
torch.onnx.export(
|
||
wrapped_model,
|
||
dummy_input,
|
||
onnx_path,
|
||
export_params=True,
|
||
opset_version=14,
|
||
input_names=['inputs'],
|
||
output_names=['predictions'],
|
||
)
|
||
print(f"✓ 简化版ONNX模型已导出: {onnx_path}")
|
||
except Exception as e2:
|
||
print(f"✗ 简化导出也失败: {e2}")
|
||
return
|
||
|
||
# ========== 步骤6: 验证ONNX模型 ==========
|
||
print("\n[6/6] 验证ONNX模型...")
|
||
|
||
try:
|
||
import onnx
|
||
import onnxruntime as ort
|
||
|
||
# 加载并检查ONNX模型
|
||
onnx_model = onnx.load(onnx_path)
|
||
onnx.checker.check_model(onnx_model)
|
||
print("✓ ONNX模型结构验证通过")
|
||
|
||
# 显示模型信息
|
||
print(f"\n模型信息:")
|
||
print(f" 输入:")
|
||
for input_tensor in onnx_model.graph.input:
|
||
dims = [d.dim_value if d.dim_value > 0 else 'dynamic' for d in input_tensor.type.tensor_type.shape.dim]
|
||
print(f" - {input_tensor.name}: {dims}")
|
||
|
||
print(f" 输出:")
|
||
for output_tensor in onnx_model.graph.output:
|
||
dims = [d.dim_value if d.dim_value > 0 else 'dynamic' for d in output_tensor.type.tensor_type.shape.dim]
|
||
print(f" - {output_tensor.name}: {dims}")
|
||
|
||
# 创建ONNX Runtime推理会话
|
||
print("\n 创建ONNX Runtime推理会话...")
|
||
ort_session = ort.InferenceSession(onnx_path)
|
||
|
||
# 准备测试输入
|
||
test_input = np.random.randn(1, 512).astype(np.float32)
|
||
test_mask = np.ones((1, 512), dtype=np.float32)
|
||
|
||
# 运行推理
|
||
print(" 运行推理测试...")
|
||
outputs = ort_session.run(
|
||
None,
|
||
{'inputs': test_input, 'masks': test_mask}
|
||
)
|
||
|
||
print(f"✓ ONNX推理成功")
|
||
print(f" 输出数量: {len(outputs)}")
|
||
for i, output in enumerate(outputs):
|
||
print(f" 输出{i}形状: {output.shape}")
|
||
|
||
except ImportError:
|
||
print("⚠️ 未安装onnx或onnxruntime，跳过验证")
|
||
print(" 安装命令: pip install onnx onnxruntime")
|
||
except Exception as e:
|
||
print(f"⚠️ ONNX验证失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
print(" 但模型文件已保存，可以尝试手动验证")
|
||
|
||
# ========== 总结 ==========
|
||
print("\n" + "="*80)
|
||
print("✓ TimesFM 2.5模型导出完成!")
|
||
print("="*80)
|
||
|
||
if os.path.exists(onnx_path):
|
||
print(f"\n保存位置: {onnx_path}")
|
||
print(f"文件大小: {os.path.getsize(onnx_path) / (1024**2):.2f} MB")
|
||
|
||
print(f"\n使用方法:")
|
||
print(f"```python")
|
||
print(f"import onnxruntime as ort")
|
||
print(f"import numpy as np")
|
||
print(f"")
|
||
print(f"# 加载模型")
|
||
print(f"session = ort.InferenceSession('{onnx_path}')")
|
||
print(f"")
|
||
print(f"# 准备输入（512小时BTC历史价格）")
|
||
print(f"btc_prices = np.array([85000, 85100, ...]) # 512个价格点")
|
||
print(f"inputs = btc_prices.reshape(1, -1).astype(np.float32)")
|
||
print(f"masks = np.ones((1, 512), dtype=np.float32)")
|
||
print(f"")
|
||
print(f"# 推理")
|
||
print(f"outputs = session.run(None, {{'inputs': inputs, 'masks': masks}})")
|
||
print(f"predictions = outputs[0] # 未来24小时价格预测")
|
||
print(f"```")
|
||
|
||
print("\n" + "="*80)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
export_timesfm_to_onnx()
|