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

This commit is contained in:
crypto_predictor committed 2026-02-08 15:05:09 +08:00
1 parent 5ca1779d7d
commit 7990312878
2 files changed
+37

No files matched your search

+37
View File
@@ -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()