transformers

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

16
import argparse
17
import json
18
import os
19

20
from tensorflow.core.protobuf.saved_model_pb2 import SavedModel
21

22

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

27
# Internal TensorFlow ops that can be safely ignored (mostly specific to a saved model)
28
INTERNAL_OPS = [
29
    "Assert",
30
    "AssignVariableOp",
31
    "EmptyTensorList",
32
    "MergeV2Checkpoints",
33
    "ReadVariableOp",
34
    "ResourceGather",
35
    "RestoreV2",
36
    "SaveV2",
37
    "ShardedFilename",
38
    "StatefulPartitionedCall",
39
    "StaticRegexFullMatch",
40
    "VarHandleOp",
41
]
42

43

44
def onnx_compliancy(saved_model_path, strict, opset):
45
    saved_model = SavedModel()
46
    onnx_ops = []
47

48
    with open(os.path.join(REPO_PATH, "utils", "tf_ops", "onnx.json")) as f:
49
        onnx_opsets = json.load(f)["opsets"]
50

51
    for i in range(1, opset + 1):
52
        onnx_ops.extend(onnx_opsets[str(i)])
53

54
    with open(saved_model_path, "rb") as f:
55
        saved_model.ParseFromString(f.read())
56

57
    model_op_names = set()
58

59
    # Iterate over every metagraph in case there is more than one (a saved model can contain multiple graphs)
60
    for meta_graph in saved_model.meta_graphs:
61
        # Add operations in the graph definition
62
        model_op_names.update(node.op for node in meta_graph.graph_def.node)
63

64
        # Go through the functions in the graph definition
65
        for func in meta_graph.graph_def.library.function:
66
            # Add operations in each function
67
            model_op_names.update(node.op for node in func.node_def)
68

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

73
    for op in model_op_names:
74
        if op not in onnx_ops and op not in INTERNAL_OPS:
75
            incompatible_ops.append(op)
76

77
    if strict and len(incompatible_ops) > 0:
78
        raise Exception(f"Found the following incompatible ops for the opset {opset}:\n" + incompatible_ops)
79
    elif len(incompatible_ops) > 0:
80
        print(f"Found the following incompatible ops for the opset {opset}:")
81
        print(*incompatible_ops, sep="\n")
82
    else:
83
        print(f"The saved model {saved_model_path} can properly be converted with ONNX.")
84

85

86
if __name__ == "__main__":
87
    parser = argparse.ArgumentParser()
88
    parser.add_argument("--saved_model_path", help="Path of the saved model to check (the .pb file).")
89
    parser.add_argument(
90
        "--opset", default=12, type=int, help="The ONNX opset against which the model has to be tested."
91
    )
92
    parser.add_argument(
93
        "--framework", choices=["onnx"], default="onnx", help="Frameworks against which to test the saved model."
94
    )
95
    parser.add_argument(
96
        "--strict", action="store_true", help="Whether make the checking strict (raise errors) or not (raise warnings)"
97
    )
98
    args = parser.parse_args()
99

100
    if args.framework == "onnx":
101
        onnx_compliancy(args.saved_model_path, args.strict, args.opset)
102