diff --git a/models/checkpoints/nbeats_btc_1h/model.onnx b/models/checkpoints/nbeats_btc_1h/model.onnx new file mode 100644 index 0000000..ab96487 Binary files /dev/null and b/models/checkpoints/nbeats_btc_1h/model.onnx differ diff --git a/models/checkpoints/ttm_r2_onnx/model.onnx b/models/checkpoints/ttm_r2_onnx/model.onnx new file mode 100644 index 0000000..adefd1c Binary files /dev/null and b/models/checkpoints/ttm_r2_onnx/model.onnx differ diff --git a/scripts/export_onnx.py b/scripts/export_onnx.py index 656f199..01ce2c5 100644 --- a/scripts/export_onnx.py +++ b/scripts/export_onnx.py @@ -1,13 +1,9 @@ -"""导出 DLinear 模型为 ONNX 格式。 - -用法: - python scripts/export_onnx.py - -输出: - models/checkpoints/dlinear_classification_2026/model.onnx -""" +"""导出 DLinear, N-BEATS 和 TTM-R2 模型为 ONNX 格式。""" import sys +import os +import torch +import numpy as np from pathlib import Path # 添加项目根目录到路径 @@ -15,23 +11,125 @@ project_root = Path(__file__).parent.parent sys.path.insert(0, str(project_root)) from models.dlinear.model import DLinear +from models.nbeats.model import NBeats +from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction +class TTMWrapper(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, x): + output = self.model(x) + return output.prediction_outputs + +def export_dlinear(): + print("\n" + "="*50) + print("导出 DLinear 模型...") + try: + checkpoint_path = project_root / "models/checkpoints/dlinear_classification_2026/best_model.pt" + output_path = checkpoint_path.parent / "model.onnx" + + if not checkpoint_path.exists(): + print(f"跳过: 找不到检查点 {checkpoint_path}") + return + + print(f"加载模型: {checkpoint_path}") + model = DLinear.load(str(checkpoint_path), device='cpu') + + print(f"模型配置: lookback={model.config.lookback}, input_dim={model.config.input_dim}") + + model.to_onnx(str(output_path)) + print(f"✓ ONNX 模型已保存到: {output_path}") + except Exception as e: + print(f"导出 DLinear 失败: {e}") + +def export_nbeats(): + print("\n" + "="*50) + print("导出 N-BEATS 模型...") + try: + checkpoint_path = project_root / "models/checkpoints/nbeats_btc_1h/best_model.pt" + output_path = checkpoint_path.parent / "model.onnx" + + if not checkpoint_path.exists(): + print(f"跳过: 找不到检查点 {checkpoint_path}") + return + + print(f"加载模型: {checkpoint_path}") + checkpoint = torch.load(str(checkpoint_path), map_location='cpu') + model = NBeats(checkpoint['config']) + model.load_state_dict(checkpoint['model_state_dict']) + model.eval() + + # N-BEATS input: (batch_size, input_size) + # input_size default is 168 + input_size = checkpoint['config'].input_size + dummy_input = torch.randn(1, input_size) + + print(f"导出 ONNX (input_size={input_size})...") + torch.onnx.export( + model, + dummy_input, + str(output_path), + export_params=True, + opset_version=14, + do_constant_folding=True, + input_names=['input'], + output_names=['forecast'], + dynamic_axes={'input': {0: 'batch_size'}, 'forecast': {0: 'batch_size'}} + ) + print(f"✓ ONNX 模型已保存到: {output_path}") + except Exception as e: + print(f"导出 N-BEATS 失败: {e}") + import traceback + traceback.print_exc() + +def export_ttm_r2(): + print("\n" + "="*50) + print("导出 TTM-R2 模型...") + try: + model_name = "ibm-granite/granite-timeseries-ttm-r2" + output_dir = project_root / "models/checkpoints/ttm_r2_onnx" + output_path = output_dir / "model.onnx" + + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"加载模型: {model_name}") + base_model = TinyTimeMixerForPrediction.from_pretrained( + model_name, + trust_remote_code=True + ) + + # Wrap model + model = TTMWrapper(base_model) + model.eval() + + # TTM input: (batch_size, context_length, num_features) + # Context length 512, features 1 + dummy_input = torch.randn(1, 512, 1) + + print(f"导出 ONNX...") + torch.onnx.export( + model, + dummy_input, + str(output_path), + export_params=True, + opset_version=14, + do_constant_folding=True, + input_names=['past_values'], + output_names=['prediction_outputs'], + dynamic_axes={'past_values': {0: 'batch_size'}, 'prediction_outputs': {0: 'batch_size'}} + ) + print(f"✓ ONNX 模型已保存到: {output_path}") + except Exception as e: + print(f"导出 TTM-R2 失败: {e}") + import traceback + traceback.print_exc() def main(): - # 模型路径 - checkpoint_path = project_root / "models/checkpoints/dlinear_classification_2026/best_model.pt" - output_path = checkpoint_path.parent / "model.onnx" - - print(f"加载模型: {checkpoint_path}") - model = DLinear.load(str(checkpoint_path), device='cpu') - - print(f"模型配置: lookback={model.config.lookback}, input_dim={model.config.input_dim}") - print(f"参数量: {model.count_parameters():,}") - - # 导出为 ONNX - model.to_onnx(str(output_path)) - print(f"✓ ONNX 模型已保存到: {output_path}") - + export_dlinear() + export_nbeats() + export_ttm_r2() if __name__ == "__main__": main()