How to Convert Tensoflow/Keras Model to the ONNX Format
2 min read4 days ago
To convert a TensorFlow/Keras model to the ONNX format, the industry-standard tool is tf2onnx.
Here is a complete, production-ready python example that builds a simple Keras model, saves it, and converts it to .onnx.
1. Install Required Libraries
You need to install TensorFlow and the official tf2onnx library.
bash
pip install tensorflow tf2onnx onnxruntimeUse code with caution.
2. Python Conversion Script
Get Dhiraj Patra’s stories in your inbox
Join Medium for free to get updates from this writer.
python
import numpy as np
import onnxruntime as ort
import tensorflow as tf
import tf2onnx# -------------------------------------------------------------
# STEP 1: Create and train a dummy Keras model
# -------------------------------------------------------------
print("Building Keras model...")
inputs = tf.keras.Input(shape=(10,), name="input_layer")
x = tf.keras.layers.Dense(32, activation="relu")(inputs)
outputs = tf.keras.layers.Dense(2, activation="softmax", name="output_layer")(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)# Dummy weights initialization
model.compile(optimizer="adam", loss="categorical_crossentropy")# -------------------------------------------------------------
# STEP 2: Define the input signature and convert to ONNX
# -------------------------------------------------------------
print("Converting model to ONNX...")# Define the shape and type of the input tensor.
# Using 'None' for the first dimension allows for dynamic batch sizes.
input_signature = [tf.TensorSpec([None, 10], tf.float32, name="input_layer")]# Perform the conversion
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature, opset=13)# Save the ONNX model to disk
onnx_path = "model.onnx"
with open(onnx_path, "wb") as f:
f.write(onnx_model.SerializeToString())print(f"Success! Model saved to {onnx_path}")# -------------------------------------------------------------
# STEP 3: Verify the ONNX model using ONNX Runtime
# -------------------------------------------------------------
print("\nVerifying ONNX model against original Keras model...")# Generate random dummy data for a test prediction
dummy_input = np.random.randn(1, 10).astype(np.float32)# Get Keras prediction
keras_pred = model.predict(dummy_input, verbose=0)# Get ONNX Runtime prediction
ort_session = ort.InferenceSession(onnx_path)
# ONNX requires a dict mapping input layer names to numpy arrays
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input}
ort_pred = ort_session.run(None, ort_inputs)[0]# Check if outputs match closely
np.testing.assert_allclose(keras_pred, ort_pred, rtol=1e-5, atol=1e-5)
print("Verification passed! Keras and ONNX outputs are identical.")
print(f"Keras output: {keras_pred}")
print(f"ONNX output: {ort_pred}")
Use code with caution.
Key Considerations for Production
- Opset Version: The
opset=13parameter defines the ONNX operator set version. Newer versions support newer neural network layers. Opset 13 to 15 is generally a safe, highly compatible choice for most edge devices. - Dynamic Axes: By passing
[None, 10]in theTensorSpec, you explicitly tell ONNX that the batch size can change at runtime (e.g., passing 1 image or 32 images at once). - Command Line Alternative: If you already have a model saved on your disk as a SavedModel folder (
my_model_dir), you don't even need to write Python code. You can convert it directly from your terminal: - bash
python -m tf2onnx.convert --saved-model my_model_dir --output model.onnx- Use code with caution.
You will find many tensorflow/keras model including my patented product application for renewable energy sector in my github repositories.
