feat: 添加 ONNX 导出功能及 DLinear 模型导出

2026-02-08 15:05:09 +08:00
2 changed files with 37 additions and 0 deletions
@@ -0,0 +1,37 @@
+"""导出 DLinear 模型为 ONNX 格式。
+
+用法:
+    python scripts/export_onnx.py
+
+输出:
+    models/checkpoints/dlinear_classification_2026/model.onnx
+"""
+
+import sys
+from pathlib import Path
+
+# 添加项目根目录到路径
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from models.dlinear.model import DLinear
+
+
+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}")
+
+
+if __name__ == "__main__":
+    main()