Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7990312878 |
No files matched your search
Binary file not shown.
@@ -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()
|
||||||
Reference in new issue
Block a user