File size: 903 Bytes
81e28a1 |
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 |
import sys
import os
import onnx
from onnx_tf.backend import prepare
def onnx_to_torch_converter(dir_name):
if not os.path.exists(dir_name):
print(f"Directory {dir_name} does not exist!")
return
onnx_model_path = os.path.join(dir_name, "onnx", "model.onnx")
if not os.path.exists(onnx_model_path):
print(f"ONNX model at {onnx_model_path} does not exist!")
return
onnx_model = onnx.load(onnx_model_path)
tf_rep = prepare(onnx_model) # prepare tf representation
tf_model_save_path = os.path.join(dir_name, "tf_model")
tf_rep.export_graph(tf_model_save_path) # export the model
print(f"PyTorch model saved at {tf_model_save_path}")
if __name__ == "__main__":
if len(sys.argv) != 2:
print("Usage: python onnx2torch.py [directory_path]")
else:
dir_name = sys.argv[1]
onnx_to_torch_converter(dir_name)
|