"""将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()