feat: Export N-BEATS and TTM-R2 models to ONNX format

scripts/export_onnx.py

@@ -1,13 +1,9 @@
-"""导出 DLinear 模型为 ONNX 格式。
+"""导出 DLinear, N-BEATS 和 TTM-R2 模型为 ONNX 格式。"""
-
-用法:
-    python scripts/export_onnx.py
-
-输出:
-    models/checkpoints/dlinear_classification_2026/model.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():
-    # 模型路径
+    export_dlinear()
-    checkpoint_path = project_root / "models/checkpoints/dlinear_classification_2026/best_model.pt"
+    export_nbeats()
-    output_path = checkpoint_path.parent / "model.onnx"
+    export_ttm_r2()
-
-    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}")
-
 
 if __name__ == "__main__":
     main()