tf2onnx Documentation
repository·main·Indexed 25 days ago
https://github.com/onnx/tensorflow-onnxA conversion tool used to migrate models from the TensorFlow ecosystem—including Keras, TF.js, and TFLite—to the ONNX format for cross-platform deployment. It provides both a command-line interface and a Python API for converting SavedModels, checkpoints, GraphDefs, and TFLite files, supporting ONNX opsets 14 through 18.
What's inside tf2onnx
- tf2onnx is a tool used to convert TensorFlow (tf-2.x), Keras, TensorFlow.js, and TFLite models into the ONNX format. Conversion can be performed using either a command-line interface or a Python API.
How tf2onnx works
mainThe
tf2onnxconversion process follows these conceptual steps:- Frozen Graph: Starts with a frozen TensorFlow graph to ensure parameters are constant.
- Protobuf Conversion: Performs a 1:1 conversion from TensorFlow protobuf to ONNX protobuf format using
tensorflow_to_onnx(). This creates an internal ONNX representation wrapped inGraphandNodeobjects. - Subgraph Rewriting: Applies graph matching to rewrite subgraphs (e.g., for
transposeorlstmops). - Individual Op Processing: Maps TensorFlow op types to processing methods via
_OPS_MAPPING. This handles broadcasting, complex compositions (likerelu6), and direct mappings. - Graph Optimization: Optimizes the functional ONNX graph by removing unnecessary ops, fusing ops, and de-duplicating constants.
- Final Processing: Performs a topological sort to ensure the graph meets ONNX requirements via
process_tf_graph().
How to find TensorFlow graph inputs and outputs
mainIf you are using
checkpointorgraphdefformats and do not know the input/output node names, use the TensorFlowsummarize_graphutility.Example usage:
summarize_graph --in_graph=path/to/your/model.pbConfigure the ONNX opset version
maintf2onnx uses the ONNX version installed on your system (or installs the latest if none is found). It supports and tests ONNX opsets 14 through 18.
By default, the resulting ONNX graph is generated using
opset-15. To specify a different opset, use the--opsetflag in the command line.--opset 15Convert Huggingface BERT models to ONNX
mainFor Natural Language Processing (NLP) tasks involving Huggingface BERT models, use the
Huggingface Bert Exampletutorial. Note that the original TensorFlow BERT model tutorial is deprecated in favor of this Huggingface approach.https://github.com/onnx/tensorflow-onnx/blob/main/tutorials/huggingface-bert.ipynbConvert TensorFlow models with custom operators to ONNX
mainWhen a TensorFlow model uses custom operators not recognized by
tf2onnxor ONNX Runtime, you must convert the operator into ONNX format. This is achieved through one of two strategies:- Decomposition: Combine existing ONNX standard operators to represent the custom logic.
- Custom ONNX Operator: Implement a new custom operator in ONNX Runtime (required if the logic cannot be decomposed into standard ops).
High-level Workflow
- Implement the TensorFlow custom operator in C++ and register it with TensorFlow (producing a
.solibrary). - Convert the operator to ONNX using a handler that either decomposes the op into standard ONNX ops or maps it to a custom ONNX op.
- Register the operator in ONNX Runtime (only necessary if using a custom ONNX op in step 2).
Convert TensorFlow Checkpoint or GraphDef to ONNX
mainIf your model is in
checkpointorgraphdefformat, you must explicitly provide the input and output node names using the--inputsand--outputsflags. Node names typically end in:0(e.g.,input0:0).If you have unknown dimensions, you can specify shapes using
[]notation, such as--inputs X:0[1,28,28,3]. Use-1for unknown dimensions.Convert TFLite models to ONNX
mainYou can convert
.tflitemodels via the command line using the--tfliteflag. Inputs and outputs do not need to be specified manually.python -m tf2onnx.convert --opset 16 --tflite tflite-file --output model.onnxConvert TensorFlow.js models to ONNX
mainUse the--tfjsflag followed by the path to the.tfjsfile. Inputs and outputs are automatically handled.How to extend tf2onnx with new operator conversions
mainTo add support for new TensorFlow operators to
tf2onnx, follow these steps:- Check existing mappings: Determine if the operator fits into an existing mapping in
_OPS_MAPPING. If it does, simply add it there. - Create a new mapping function: If the operator requires specialized processing during conversion, implement a new mapping function.
- Use graph re-writes: For complex patterns where a single TensorFlow op is composed of multiple ops, consider using a graph re-write approach for better stability.
- Add unit tests: Implement tests in
tests/test_backend.py. The standard pattern is to create a TensorFlow graph, capture its output, convert it to ONNX, run it against an ONNX backend, and compare the results. - Test with pre-trained models: If available, add models using the new operator to
test/run_pretrained_models.pyto ensure end-to-end correctness.
- Check existing mappings: Determine if the operator fits into an existing mapping in
Prerequisites for tf2onnx
mainBefore using
tf2onnx, ensure you have the following installed:- TensorFlow: Supports
tf-2.x. It is recommended to usetf-2.13or better.pip install tensorflow
2. **ONNX**: Requires `onnx-1.9` or better (will be installed/upgraded automatically). 3. **(Optional) Runtime**: If you intend to run tests, install an ONNX runtime such as `onnxruntime`. ```bash pip install onnxruntime- TensorFlow: Supports
Convert TFLite models to ONNX
mainTo convert a TensorFlow Lite (
.tflite) model to ONNX, use the workflow demonstrated in themobiledettutorial. This is specifically designed for handling the TFLite format.https://github.com/onnx/tensorflow-onnx/blob/main/tutorials/mobiledet-tflite.ipynb