135 lines
4.4 KiB
Python
135 lines
4.4 KiB
Python
"""导出 DLinear, N-BEATS 和 TTM-R2 模型为 ONNX 格式。"""
|
|
|
|
import sys
|
|
import os
|
|
import torch
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
# 添加项目根目录到路径
|
|
project_root = Path(__file__).parent.parent
|
|
sys.path.insert(0, str(project_root))
|
|
|
|
from models.dlinear.model import DLinear
|
|
from models.nbeats.model import NBeats
|
|
from tsfm_public.toolkit.get_model import get_model
|
|
|
|
|
|
def export_dlinear():
|
|
print("\n" + "="*50)
|
|
print("导出 DLinear 模型...")
|
|
try:
|
|
checkpoint_path = project_root / "models/checkpoints/dlinear_classification_2026/best_model.pt"
|
|
output_path = checkpoint_path.parent / "model.onnx"
|
|
|
|
if not checkpoint_path.exists():
|
|
print(f"跳过: 找不到检查点 {checkpoint_path}")
|
|
return
|
|
|
|
print(f"加载模型: {checkpoint_path}")
|
|
model = DLinear.load(str(checkpoint_path), device='cpu')
|
|
|
|
print(f"模型配置: lookback={model.config.lookback}, input_dim={model.config.input_dim}")
|
|
|
|
model.to_onnx(str(output_path))
|
|
print(f"✓ ONNX 模型已保存到: {output_path}")
|
|
except Exception as e:
|
|
print(f"导出 DLinear 失败: {e}")
|
|
|
|
def export_nbeats():
|
|
print("\n" + "="*50)
|
|
print("导出 N-BEATS 模型...")
|
|
try:
|
|
checkpoint_path = project_root / "models/checkpoints/nbeats_btc_1h/best_model.pt"
|
|
output_path = checkpoint_path.parent / "model.onnx"
|
|
|
|
if not checkpoint_path.exists():
|
|
print(f"跳过: 找不到检查点 {checkpoint_path}")
|
|
return
|
|
|
|
print(f"加载模型: {checkpoint_path}")
|
|
checkpoint = torch.load(str(checkpoint_path), map_location='cpu')
|
|
model = NBeats(checkpoint['config'])
|
|
model.eval()
|
|
|
|
# N-BEATS input: (batch_size, input_size)
|
|
# input_size default is 168
|
|
input_size = checkpoint['config'].input_size
|
|
dummy_input = torch.randn(1, input_size)
|
|
|
|
print(f"导出 ONNX (input_size={input_size})...")
|
|
torch.onnx.export(
|
|
model,
|
|
dummy_input,
|
|
str(output_path),
|
|
export_params=True,
|
|
opset_version=17,
|
|
do_constant_folding=True,
|
|
input_names=['input'],
|
|
output_names=['forecast'],
|
|
dynamic_axes={'input': {0: 'batch_size'}, 'forecast': {0: 'batch_size'}}
|
|
)
|
|
print(f"✓ ONNX 模型已保存到: {output_path}")
|
|
except Exception as e:
|
|
print(f"导出 N-BEATS 失败: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
def export_ttm_r2():
|
|
print("\n" + "="*50)
|
|
print("导出 TTM-R2 模型...")
|
|
try:
|
|
model_name = "ibm-granite/granite-timeseries-ttm-r2"
|
|
output_dir = project_root / "models/checkpoints/ttm_r2_onnx"
|
|
output_path = output_dir / "model.onnx"
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
print(f"加载模型: {model_name}")
|
|
base_model = get_model(
|
|
model_path=model_name,
|
|
context_length=512,
|
|
prediction_length=96,
|
|
)
|
|
base_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)
|
|
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(
|
|
base_model,
|
|
dummy_input,
|
|
str(output_path),
|
|
export_params=True,
|
|
opset_version=17, # Revert to 14 as it is standard
|
|
do_constant_folding=True,
|
|
input_names=['past_values'],
|
|
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:
|
|
print(f"导出 TTM-R2 失败: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
def main():
|
|
export_dlinear()
|
|
export_nbeats()
|
|
#export_ttm_r2()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|