62 lines
1.4 KiB
Python
62 lines
1.4 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Convert TensorFlow SavedModel to ONNX format."""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
description="Convert TensorFlow SavedModel to ONNX format"
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--saved-model",
|
||
|
|
type=Path,
|
||
|
|
default=Path("model/saved_model"),
|
||
|
|
help="Path to the TensorFlow SavedModel directory",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--output",
|
||
|
|
type=Path,
|
||
|
|
default=Path("model/model.onnx"),
|
||
|
|
help="Path to the output ONNX model file",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--opset",
|
||
|
|
type=int,
|
||
|
|
default=13,
|
||
|
|
help="ONNX opset version to target",
|
||
|
|
)
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
# Check that saved model exists
|
||
|
|
if not args.saved_model.exists():
|
||
|
|
raise FileNotFoundError(f"SavedModel not found at {args.saved_model}")
|
||
|
|
|
||
|
|
# Create output parent directory if needed
|
||
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
# Run tf2onnx conversion
|
||
|
|
cmd = [
|
||
|
|
sys.executable,
|
||
|
|
"-m",
|
||
|
|
"tf2onnx.convert",
|
||
|
|
"--saved-model",
|
||
|
|
str(args.saved_model),
|
||
|
|
"--output",
|
||
|
|
str(args.output),
|
||
|
|
"--opset",
|
||
|
|
str(args.opset),
|
||
|
|
]
|
||
|
|
|
||
|
|
subprocess.run(cmd, check=True)
|
||
|
|
|
||
|
|
print(f"ONNX model saved to {args.output}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|