Implementation:Microsoft LoRA Check TF Ops
Appearance
Overview
The check_tf_ops.py utility checks whether the TensorFlow operations used in a SavedModel are compatible with a specified ONNX opset, enabling validation before ONNX conversion.
Description
When exporting TensorFlow SavedModels to ONNX format, certain TensorFlow operations may not have corresponding ONNX operator implementations, causing conversion failures. This script preemptively identifies incompatible operations by:
- Loading the ONNX op registry: Reads
utils/tf_ops/onnx.json, which contains a dictionary mapping ONNX opset versions (1 through N) to lists of supported TF operation names. Accumulates all ops from opset 1 through the target opset.
- Parsing the SavedModel protobuf: Loads the
.pbfile using TensorFlow'sSavedModelprotobuf class. Iterates over all meta-graphs, collecting unique operation names from:- Top-level graph definition nodes (
meta_graph.graph_def.node) - Function library nodes (
meta_graph.graph_def.library.function[].node_def)
- Top-level graph definition nodes (
- Comparing against ONNX ops: Filters out known internal TensorFlow ops that are specific to SavedModel infrastructure (listed in
INTERNAL_OPS:Assert,AssignVariableOp,EmptyTensorList,MergeV2Checkpoints,ReadVariableOp,ResourceGather,RestoreV2,SaveV2,ShardedFilename,StatefulPartitionedCall,StaticRegexFullMatch,VarHandleOp). Reports remaining ops not in the ONNX opset.
- Reporting: In strict mode, raises an exception on incompatible ops. In non-strict mode, prints warnings. If all ops are compatible, prints a success message.
Usage
Use this utility when:
- Validating a TensorFlow SavedModel can be converted to ONNX before running the actual conversion.
- Debugging ONNX conversion failures by identifying which TF ops are unsupported.
- Testing model ONNX compatibility as part of a CI pipeline at a specific opset version.
Code Reference
Source Location
examples/NLU/utils/check_tf_ops.py (101 lines)
Signature
INTERNAL_OPS = [
"Assert", "AssignVariableOp", "EmptyTensorList", "MergeV2Checkpoints",
"ReadVariableOp", "ResourceGather", "RestoreV2", "SaveV2",
"ShardedFilename", "StatefulPartitionedCall", "StaticRegexFullMatch", "VarHandleOp",
]
def onnx_compliancy(saved_model_path: str, strict: bool, opset: int) -> None: ...
Import / CLI Usage
# Check ONNX compatibility at opset 12
python utils/check_tf_ops.py \
--saved_model_path ./saved_model/saved_model.pb \
--opset 12
# Strict mode (raises exception on incompatibility)
python utils/check_tf_ops.py \
--saved_model_path ./saved_model/saved_model.pb \
--opset 12 \
--strict
I/O Contract
Inputs
| Input | Type | Description |
|---|---|---|
--saved_model_path |
str | Path to the TensorFlow SavedModel .pb file to check
|
--opset |
int | Target ONNX opset version to validate against. Default: 12 |
--framework |
str | Framework to test against (currently only onnx is supported). Default: onnx
|
--strict |
flag | If set, raises an exception on incompatible ops; otherwise prints warnings |
utils/tf_ops/onnx.json |
JSON file | Registry mapping ONNX opset versions to lists of supported TensorFlow operations |
Outputs
| Output | Type | Description |
|---|---|---|
| Success message | stdout | "The saved model {path} can properly be converted with ONNX." when all ops are compatible
|
| Warning list | stdout | List of incompatible ops (non-strict mode) |
| Exception | Exception | Raised with incompatible op names (strict mode) |
Usage Examples
# Basic compatibility check at default opset 12
python utils/check_tf_ops.py \
--saved_model_path ./exported_model/saved_model.pb
# Output (success):
# The saved model ./exported_model/saved_model.pb can properly be converted with ONNX.
# Output (incompatible):
# Found the following incompatible ops for the opset 12:
# EuclideanNorm
# FusedBatchNormV3
# Strict mode - raises exception on failure
python utils/check_tf_ops.py \
--saved_model_path ./exported_model/saved_model.pb \
--opset 13 \
--strict
# Check at a higher opset version for wider compatibility
python utils/check_tf_ops.py \
--saved_model_path ./exported_model/saved_model.pb \
--opset 15
Related Pages
Page Connections
Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment