Upload ai/utils/export_onnx.py with huggingface_hub

ai/utils/export_onnx.py

@@ -0,0 +1,41 @@
+import os
+import sys
+
+import torch
+
+# Add project root to path
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+from ai.models.training_config import INPUT_SIZE, POLICY_SIZE
+from ai.training.train import AlphaNet
+
+
+def export_to_onnx(model_path, output_path):
+    device = torch.device("cpu")
+    checkpoint = torch.load(model_path, map_location=device)
+    state_dict = (
+        checkpoint["model_state"] if isinstance(checkpoint, dict) and "model_state" in checkpoint else checkpoint
+    )
+
+    model = AlphaNet(policy_size=POLICY_SIZE).to(device)
+    model.load_state_dict(state_dict)
+    model.eval()
+
+    dummy_input = torch.randn(1, INPUT_SIZE)
+
+    torch.onnx.export(
+        model,
+        dummy_input,
+        output_path,
+        export_params=True,
+        opset_version=14,
+        do_constant_folding=True,
+        input_names=["input"],
+        output_names=["policy", "value"],
+        dynamic_axes={"input": {0: "batch_size"}, "policy": {0: "batch_size"}, "value": {0: "batch_size"}},
+    )
+    print(f"Model successfully exported to {output_path}")
+
+
+if __name__ == "__main__":
+    export_to_onnx("ai/models/alphanet_best.pt", "ai/models/alphanet.onnx")