diff --git a/Machine_Learning/convert_onnx.py b/Machine_Learning/convert_onnx.py new file mode 100644 index 0000000..b2a0eb2 --- /dev/null +++ b/Machine_Learning/convert_onnx.py @@ -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() diff --git a/Machine_Learning/requirements.txt b/Machine_Learning/requirements.txt index 4094d4d..5f539d9 100644 --- a/Machine_Learning/requirements.txt +++ b/Machine_Learning/requirements.txt @@ -2,3 +2,6 @@ tensorflow>=2.13.0 tensorflowjs>=4.10.0 jupyter>=1.0.0 ipykernel>=6.25.0 +tf2onnx>=1.16.1 +onnx>=1.16.0 +onnxruntime>=1.17.0