Spaces:
Sleeping
Sleeping
File size: 2,087 Bytes
37004b0 |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import argparse
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_model(model_path, output_path):
device = torch.device("cpu") # ONNX export usually done on CPU to be safe
print(f"Loading checkpoint from {model_path}...")
checkpoint = torch.load(model_path, map_location=device)
# Handle dictionary checkpoint
if isinstance(checkpoint, dict) and "model_state" in checkpoint:
state_dict = checkpoint["model_state"]
else:
state_dict = checkpoint
# Detect policy size from weights to support different versions
if "policy_head_fc.bias" in state_dict:
detected_size = state_dict["policy_head_fc.bias"].shape[0]
print(f"Detected Policy Size: {detected_size}")
else:
detected_size = POLICY_SIZE
print(f"Using Default Policy Size: {detected_size}")
model = AlphaNet(policy_size=detected_size).to(device)
model.load_state_dict(state_dict)
model.eval()
dummy_input = torch.randn(1, INPUT_SIZE, device=device)
print(f"Exporting to {output_path}...")
torch.onnx.export(
model,
dummy_input,
output_path,
export_params=True,
opset_version=17, # Use modern opset for better compatibility
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("Export complete.")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, required=True, help="Path to .pt checkpoint")
parser.add_argument("--output", type=str, required=True, help="Path to .onnx output")
args = parser.parse_args()
export_model(args.model, args.output)
|