feat: Update ONNX export script and add TTM-R2 ONNX model

scripts/export_onnx.py

@@ -12,16 +12,8 @@ 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
+from tsfm_public.toolkit.get_model import get_model
 
-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)
@@ -58,7 +50,6 @@ def export_nbeats():
         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)
@@ -72,7 +63,7 @@ def export_nbeats():
             dummy_input,
             str(output_path),
             export_params=True,
-            opset_version=14,
+            opset_version=17,
             do_constant_folding=True,
             input_names=['input'],
             output_names=['forecast'],
@@ -95,30 +86,38 @@ def export_ttm_r2():
         output_dir.mkdir(parents=True, exist_ok=True)
         
         print(f"加载模型: {model_name}")
-        base_model = TinyTimeMixerForPrediction.from_pretrained(
-            model_name,
-            trust_remote_code=True
+        base_model = get_model(
+            model_path=model_name,
+            context_length=512,
+            prediction_length=96,
         )
+        base_model.eval()
         
-        # Wrap model
-        model = TTMWrapper(base_model)
-        model.eval()
+        print(f"TTM Config: {base_model.config}")
         
         # TTM input: (batch_size, context_length, num_features)
         # Context length 512, features 1
         dummy_input = torch.randn(1, 512, 1)
-        
-        print(f"导出 ONNX...")
+        with torch.no_grad():
+            output = base_model(dummy_input)
+            if hasattr(output, 'prediction_outputs'):
+                test_output = output.prediction_outputs
+            else:
+                test_output = output
+            print(f"✓ 输出形状: {test_output.shape}")
+        print(f"导出 ONNX (固定 Batch Size=1)...")
         torch.onnx.export(
-            model,
+            base_model,
             dummy_input,
             str(output_path),
             export_params=True,
-            opset_version=14,
+            opset_version=17,  # Revert to 14 as it is standard
             do_constant_folding=True,
             input_names=['past_values'],
-            output_names=['prediction_outputs'],
-            dynamic_axes={'past_values': {0: 'batch_size'}, 'prediction_outputs': {0: 'batch_size'}}
+            output_names=['predictions'],
+            # Remove dynamic axes to fix potential shape inference issues
+            dynamic_axes={'past_values': {0: 'batch_size'}, 'predictions': {0: 'batch_size'}},
+            verbose=False
         )
         print(f"✓ ONNX 模型已保存到: {output_path}")
     except Exception as e:
@@ -129,7 +128,7 @@ def export_ttm_r2():
 def main():
     export_dlinear()
     export_nbeats()
-    export_ttm_r2()
+    #export_ttm_r2()
 
 if __name__ == "__main__":
     main()