tf2onnx Documentation

repository·main·Indexed 25 days ago

https://github.com/onnx/tensorflow-onnx

A 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.

Tokens
20K
Snippets
39
Records
111
Agent score
81%

What's inside tf2onnx

  1. Convert TensorFlow models to ONNX with tf2onnx

    main
    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.
  2. How tf2onnx works

    main

    The tf2onnx conversion process follows these conceptual steps:

    1. Frozen Graph: Starts with a frozen TensorFlow graph to ensure parameters are constant.
    2. 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 in Graph and Node objects.
    3. Subgraph Rewriting: Applies graph matching to rewrite subgraphs (e.g., for transpose or lstm ops).
    4. Individual Op Processing: Maps TensorFlow op types to processing methods via _OPS_MAPPING. This handles broadcasting, complex compositions (like relu6), and direct mappings.
    5. Graph Optimization: Optimizes the functional ONNX graph by removing unnecessary ops, fusing ops, and de-duplicating constants.
    6. Final Processing: Performs a topological sort to ensure the graph meets ONNX requirements via process_tf_graph().
  3. How to find TensorFlow graph inputs and outputs

    main

    If you are using checkpoint or graphdef formats and do not know the input/output node names, use the TensorFlow summarize_graph utility.

    Example usage:

    summarize_graph --in_graph=path/to/your/model.pb
  4. Configure the ONNX opset version

    main

    tf2onnx 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 --opset flag in the command line.

    --opset 15
  5. Convert Huggingface BERT models to ONNX

    main

    For Natural Language Processing (NLP) tasks involving Huggingface BERT models, use the Huggingface Bert Example tutorial. 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.ipynb
  6. Convert TensorFlow models with custom operators to ONNX

    main

    When a TensorFlow model uses custom operators not recognized by tf2onnx or ONNX Runtime, you must convert the operator into ONNX format. This is achieved through one of two strategies:

    1. Decomposition: Combine existing ONNX standard operators to represent the custom logic.
    2. Custom ONNX Operator: Implement a new custom operator in ONNX Runtime (required if the logic cannot be decomposed into standard ops).

    High-level Workflow

    1. Implement the TensorFlow custom operator in C++ and register it with TensorFlow (producing a .so library).
    2. 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.
    3. Register the operator in ONNX Runtime (only necessary if using a custom ONNX op in step 2).
  7. Convert TensorFlow Checkpoint or GraphDef to ONNX

    main

    If your model is in checkpoint or graphdef format, you must explicitly provide the input and output node names using the --inputs and --outputs flags. 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 -1 for unknown dimensions.

  8. Convert TFLite models to ONNX

    main

    You can convert .tflite models via the command line using the --tflite flag. Inputs and outputs do not need to be specified manually.

    python -m tf2onnx.convert --opset 16 --tflite tflite-file --output model.onnx
  9. How to extend tf2onnx with new operator conversions

    main

    To add support for new TensorFlow operators to tf2onnx, follow these steps:

    1. Check existing mappings: Determine if the operator fits into an existing mapping in _OPS_MAPPING. If it does, simply add it there.
    2. Create a new mapping function: If the operator requires specialized processing during conversion, implement a new mapping function.
    3. 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.
    4. 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.
    5. Test with pre-trained models: If available, add models using the new operator to test/run_pretrained_models.py to ensure end-to-end correctness.
  10. Prerequisites for tf2onnx

    main

    Before using tf2onnx, ensure you have the following installed:

    1. TensorFlow: Supports tf-2.x. It is recommended to use tf-2.13 or 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
  11. Convert TFLite models to ONNX

    main

    To convert a TensorFlow Lite (.tflite) model to ONNX, use the workflow demonstrated in the mobiledet tutorial. This is specifically designed for handling the TFLite format.

    https://github.com/onnx/tensorflow-onnx/blob/main/tutorials/mobiledet-tflite.ipynb