tf2onnx_keras_model_conversion_quickstart.py

1import tensorflow as tf
2import tf2onnx
3import onnx
4
5# 1. Create a simple TensorFlow Keras model
6model = tf.keras.Sequential([
7    tf.keras.layers.Dense(10, input_shape=(5,), activation='relu'),
8    tf.keras.layers.Dense(1)
9])
10
11# 2. Define the input signature (required for conversion)
12# This specifies the shape and type of the input tensor
13spec = (tf.TensorSpec((None, 5), tf.float32, name="input"),)
14
15# 3. Convert the model to ONNX
16# output_path is optional; if provided, it saves the file directly
17model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13)
18
19# 4. Save the ONNX model to a file
20output_path = "model.onnx"
21onnx.save_model(model_proto, output_path)
22
23print(f"Model has been converted and saved to {output_path}")