refactor: 清理测试文件和2019相关代码
- 删除测试脚本: 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分类模型
This commit is contained in:
crypto_predictor committed
1 parent
c74f4f6949
commit
1bd9d25f60
33 files changed
+6557
-233
No files matched your search
@@ -0,0 +1,295 @@
|
||||
"""将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()
|
||||
Reference in new issue
Block a user