diff --git a/models/checkpoints/nbeats_btc_1h/model.onnx b/models/checkpoints/nbeats_btc_1h/model.onnx index ab96487..3a9e989 100644 Binary files a/models/checkpoints/nbeats_btc_1h/model.onnx and b/models/checkpoints/nbeats_btc_1h/model.onnx differ diff --git a/models/checkpoints/ttm_r2_onnx/model.onnx b/models/checkpoints/ttm_r2_onnx/model.onnx index adefd1c..b6ecc21 100644 Binary files a/models/checkpoints/ttm_r2_onnx/model.onnx and b/models/checkpoints/ttm_r2_onnx/model.onnx differ diff --git a/models/checkpoints/ttm_r2_onnx/ttm_r2_crypto_channel_mixing.onnx b/models/checkpoints/ttm_r2_onnx/ttm_r2_crypto_channel_mixing.onnx new file mode 100644 index 0000000..c02cb4c Binary files /dev/null and b/models/checkpoints/ttm_r2_onnx/ttm_r2_crypto_channel_mixing.onnx differ diff --git a/models/checkpoints/ttm_r2_onnx/ttm_r2_crypto_channel_mixing.onnx.data b/models/checkpoints/ttm_r2_onnx/ttm_r2_crypto_channel_mixing.onnx.data new file mode 100644 index 0000000..b89e86e Binary files /dev/null and b/models/checkpoints/ttm_r2_onnx/ttm_r2_crypto_channel_mixing.onnx.data differ diff --git a/scripts/export_onnx.py b/scripts/export_onnx.py index 01ce2c5..2969ae9 100644 --- a/scripts/export_onnx.py +++ b/scripts/export_onnx.py @@ -12,16 +12,8 @@ sys.path.insert(0, str(project_root)) from models.dlinear.model import DLinear from models.nbeats.model import NBeats -from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction +from tsfm_public.toolkit.get_model import get_model -class TTMWrapper(torch.nn.Module): - def __init__(self, model): - super().__init__() - self.model = model - - def forward(self, x): - output = self.model(x) - return output.prediction_outputs def export_dlinear(): print("\n" + "="*50) @@ -58,7 +50,6 @@ def export_nbeats(): print(f"加载模型: {checkpoint_path}") checkpoint = torch.load(str(checkpoint_path), map_location='cpu') model = NBeats(checkpoint['config']) - model.load_state_dict(checkpoint['model_state_dict']) model.eval() # N-BEATS input: (batch_size, input_size) @@ -72,7 +63,7 @@ def export_nbeats(): dummy_input, str(output_path), export_params=True, - opset_version=14, + opset_version=17, do_constant_folding=True, input_names=['input'], output_names=['forecast'], @@ -95,30 +86,38 @@ def export_ttm_r2(): output_dir.mkdir(parents=True, exist_ok=True) print(f"加载模型: {model_name}") - base_model = TinyTimeMixerForPrediction.from_pretrained( - model_name, - trust_remote_code=True + base_model = get_model( + model_path=model_name, + context_length=512, + prediction_length=96, ) + base_model.eval() - # Wrap model - model = TTMWrapper(base_model) - 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) - - print(f"导出 ONNX...") + 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( - model, + base_model, dummy_input, str(output_path), export_params=True, - opset_version=14, + opset_version=17, # Revert to 14 as it is standard do_constant_folding=True, input_names=['past_values'], - output_names=['prediction_outputs'], - dynamic_axes={'past_values': {0: 'batch_size'}, 'prediction_outputs': {0: 'batch_size'}} + 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: @@ -129,7 +128,7 @@ def export_ttm_r2(): def main(): export_dlinear() export_nbeats() - export_ttm_r2() + #export_ttm_r2() if __name__ == "__main__": main()