check_tf_ops.py

# coding=utf-8
# Copyright 2020 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

importargparse
importjson
importos

fromtensorflow.core.protobuf.saved_model_pb2importSavedModel


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_copies.py
REPO_PATH="."

# Internal TensorFlow ops that can be safely ignored (mostly specific to a saved model)
INTERNAL_OPS= [
"Assert",
"AssignVariableOp",
"EmptyTensorList",
"MergeV2Checkpoints",
"ReadVariableOp",
"ResourceGather",
"RestoreV2",
"SaveV2",
"ShardedFilename",
"StatefulPartitionedCall",
"StaticRegexFullMatch",
"VarHandleOp",
]


defonnx_compliancy(saved_model_path, strict, opset):
saved_model=SavedModel()
onnx_ops= []

withopen(os.path.join(REPO_PATH, "utils", "tf_ops", "onnx.json")) asf:
onnx_opsets=json.load(f)["opsets"]

foriinrange(1, opset+1):
onnx_ops.extend(onnx_opsets[str(i)])

withopen(saved_model_path, "rb") asf:
saved_model.ParseFromString(f.read())

model_op_names=set()

# Iterate over every metagraph in case there is more than one (a saved model can contain multiple graphs)
formeta_graphinsaved_model.meta_graphs:
# Add operations in the graph definition
model_op_names.update(node.opfornodeinmeta_graph.graph_def.node)

# Go through the functions in the graph definition
forfuncinmeta_graph.graph_def.library.function:
# Add operations in each function
model_op_names.update(node.opfornodeinfunc.node_def)

# Convert to list, sorted if you want
model_op_names=sorted(model_op_names)
incompatible_ops= []

foropinmodel_op_names:
ifopnotinonnx_opsandopnotinINTERNAL_OPS:
incompatible_ops.append(op)

ifstrictandlen(incompatible_ops) >0:
raiseException(f"Found the following incompatible ops for the opset {opset}:\n"+incompatible_ops)
eliflen(incompatible_ops) >0:
print(f"Found the following incompatible ops for the opset {opset}:")
print(*incompatible_ops, sep="\n")
else:
print(f"The saved model {saved_model_path} can properly be converted with ONNX.")


if__name__=="__main__":
parser=argparse.ArgumentParser()
parser.add_argument("--saved_model_path", help="Path of the saved model to check (the .pb file).")
parser.add_argument(
"--opset", default=12, type=int, help="The ONNX opset against which the model has to be tested."
 )
parser.add_argument(
"--framework", choices=["onnx"], default="onnx", help="Frameworks against which to test the saved model."
 )
parser.add_argument(
"--strict", action="store_true", help="Whether make the checking strict (raise errors) or not (raise warnings)"
 )
args=parser.parse_args()

ifargs.framework=="onnx":
onnx_compliancy(args.saved_model_path, args.strict, args.opset)