NautilusTrainer / debug_onnx_export.py
gionuibk's picture
Deploy NautilusTrainer with Ray RLlib - 1765520777.5218828
db75a76 verified
raw
history blame contribute delete
923 Bytes
import torch
import torch.nn as nn
from models.deeplob import DeepLOB
import os
print(f"Torch Version: {torch.__version__}")
try:
import onnx
print(f"ONNX Version: {onnx.__version__}")
except ImportError:
print("ONNX not installed")
try:
print("Initializing DeepLOB...")
model = DeepLOB(y_len=3)
model.eval()
# Dummy input: (Batch=1, Channels=2, Time=100, Features=40)
dummy = torch.randn(1, 2, 100, 40)
output_path = "debug_deeplob.onnx"
print("Exporting to ONNX...")
torch.onnx.export(
model, dummy, output_path,
input_names=['input'], output_names=['output'],
opset_version=12,
do_constant_folding=True,
verbose=True # Print compilation steps
)
print(f"✅ Export Success: {output_path}")
except Exception as e:
print(f"❌ Export Failed:")
print(e)
import traceback
traceback.print_exc()