feat: add ONNX conversion script

This commit is contained in:
Asep Haryana Saputra
2026-05-23 09:57:19 +00:00
parent 3d3f1d14fa
commit 4d82bbae34
2 changed files with 64 additions and 0 deletions
+61
View File
@@ -0,0 +1,61 @@
#!/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()