38 lines
960 B
Python
38 lines
960 B
Python
"""导出 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()
|