import onnxruntime
import torch
from nets import get_models

def pth2onnx(load='models/best_model.pth', out="models/best_model.onnx"):
    """
    Args:
        load: the pth's path
        out: the onnx's path
    """
    model = get_models("cpu", load=load)
    model.eval()
    x = torch.randn(1, 3, 224, 224, requires_grad=True)
    torch.onnx.export(model, x, out,
                    input_names=["input"], output_names=["output"],
                    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
                    opset_version=12)
    
if __name__ == "__main__":
    pth2onnx()