TensorFlow to ONNX Converter

1.17.0 · active · verified Wed Apr 15

tf2onnx is a Python library that enables the conversion of TensorFlow models into the ONNX (Open Neural Network Exchange) format. This allows users to deploy models trained in TensorFlow to various ONNX-compatible runtimes and hardware accelerators. The current version is 1.17.0, with minor releases typically occurring every 1-3 months to keep pace with TensorFlow and ONNX updates.

Warnings

breaking TensorFlow Version Compatibility: `tf2onnx` closely tracks TensorFlow releases. Older `tf2onnx` versions may not support newer TensorFlow features, and newer `tf2onnx` versions might drop support for older TensorFlow versions. This often leads to conversion failures or incorrect model behavior if not aligned.
Fix: Ensure your `tf2onnx` version is compatible with your installed `tensorflow` version. Check the `tf2onnx` release notes for specific TensorFlow version support. Upgrade both `tf2onnx` and `tensorflow` in tandem if encountering issues.
gotcha Protobuf Version Conflicts: `tf2onnx` (and its dependencies like `onnx` and `tensorflow`) have strict and often conflicting requirements for the `protobuf` library. Installing incorrect `protobuf` versions can lead to 'module not found' errors, serialization issues, or obscure runtime crashes.
Fix: After installing `tensorflow` and `onnx`, let `pip` resolve `protobuf` dependencies. If conflicts arise, consider using a virtual environment or explicitly pinning `protobuf` to a version known to work with your specific `tensorflow` and `onnx` versions (e.g., `pip install protobuf==3.20.3`).
gotcha Custom TensorFlow Operations: TensorFlow models containing custom operations, layers, or complex control flow that do not have direct equivalents in the ONNX operator set will fail conversion. `tf2onnx` provides some fallback mechanisms but they are not exhaustive.
Fix: Simplify your TensorFlow model to use standard operations where possible. For truly custom ops, you may need to implement a custom ONNX operator, or perform graph surgery (e.g., separate parts of the model for conversion).
gotcha Input/Output Signature Mismatch: Incorrectly specifying input signatures (`input_signature`) or output names (`output_names`) during conversion, especially with `from_saved_model` or `from_graph_def`, can lead to truncated models, missing outputs, or graph extraction errors.
Fix: Always explicitly define `input_signature` and `output_names` when converting. Use tools like `saved_model_cli show` to inspect your TensorFlow SavedModel's input/output signatures and ensure they match the conversion parameters exactly.

Install

pip install tf2onnx Basic installation
pip install tf2onnx tensorflow onnx onnxruntime With common dependencies

Imports

convert
```
import tf2onnx.convert
```
The 'convert' module contains the main conversion functions like from_keras, from_saved_model.

Quickstart

This quickstart demonstrates converting a simple Keras sequential model (saved as a TensorFlow SavedModel) to the ONNX format using `tf2onnx.convert.from_saved_model`. It includes saving the Keras model, performing the conversion, saving the ONNX output, and optionally verifying the ONNX model's validity.

import tensorflow as tf
import tf2onnx
from onnx.checker import check_model
import shutil
import os

# 1. Create a simple Keras model
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(10, input_shape=(784,), activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 2. Save the Keras model in the TensorFlow SavedModel format
# This is the recommended approach for converting TF2 models.
tf.saved_model.save(model, "my_keras_model")

# 3. Convert the SavedModel to ONNX
# Specify input_signature and output_names for robust conversion.
onnx_model_proto, _ = tf2onnx.convert.from_saved_model(
    "my_keras_model",
    input_signature=[tf.TensorSpec([None, 784], tf.float32, name="input_0")],
    output_names=["output_0"]
)

# 4. Save the ONNX model to a file
with open("model.onnx", "wb") as f:
    f.write(onnx_model_proto.SerializeToString())

print("ONNX model converted and saved as model.onnx")

# Optional: Verify the ONNX model structure
try:
    check_model(onnx_model_proto)
    print("ONNX model is valid.")
except Exception as e:
    print(f"ONNX model validation failed: {e}")

# Clean up created files/directories
shutil.rmtree("my_keras_model")
os.remove("model.onnx")

view raw JSON →