67 lines
1.6 KiB
Python
67 lines
1.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert TensorFlow SavedModel to ONNX format."""
|
|
|
|
import argparse
|
|
import logging
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|
|
|
|
|
def main():
|
|
"""Convert a TensorFlow SavedModel to ONNX format using tf2onnx."""
|
|
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()
|
|
|
|
if not args.saved_model.exists():
|
|
raise FileNotFoundError(f"SavedModel not found at {args.saved_model}")
|
|
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
cmd = [
|
|
sys.executable,
|
|
"-m",
|
|
"tf2onnx.convert",
|
|
"--saved-model",
|
|
str(args.saved_model),
|
|
"--output",
|
|
str(args.output),
|
|
"--opset",
|
|
str(args.opset),
|
|
]
|
|
|
|
try:
|
|
subprocess.run(cmd, check=True)
|
|
except subprocess.CalledProcessError as e:
|
|
logging.error(f"tf2onnx conversion failed with exit code {e.returncode}")
|
|
raise
|
|
|
|
logging.info(f"ONNX model saved to {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|