Spaces:
Running
Running
File size: 1,259 Bytes
1f19911 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
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")
|