1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
|
#! /usr/bin/env python3
# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
import onnx
from tf2onnx import convert
from tensorflow.python.tools import saved_model_utils
def find(seq, test, default=None):
return next((x for x in seq if test(x)), default)
def index_of(seq, elem):
for i in range(len(seq)):
if seq[i] == elem:
return i
def has_initializer(onnx_model, name):
return find(onnx_model.graph.initializer, lambda i: i.name == name) != None
def has_equivalent_shape(onnx_tensor_shape, tensorflow_tensor_shape):
onnx_dims = onnx_tensor_shape.dim
tf_dims = tensorflow_tensor_shape.dim
if len(onnx_dims) != len(tf_dims):
return False
for i in range(len(onnx_dims)):
onnx_dim_size = onnx_dims[i].dim_value
tf_dim_size = tf_dims[i].size
if onnx_dim_size == 0 and tf_dim_size == -1:
continue
if onnx_dim_size != tf_dim_size:
return False
return True
def find_by_type(seq, tensor):
return [ item for item in seq if has_equivalent_shape(item.type.tensor_type.shape, tensor.tensor_shape) ]
def rename_output(onnx_model, signature_name, signature_tensor):
signature_node_name = signature_tensor.name
graph_output = find(onnx_model.graph.output, lambda output: output.name == signature_node_name)
if graph_output is None:
print("TensorFlow signature output '{}' references node '{}' which was not found. Trying to find equivalent output.".format(signature_name, signature_node_name))
candidates = find_by_type(onnx_model.graph.output, signature_tensor)
if len(candidates) == 0:
print("Could not find equivalent output for '{}'. Unable to rename this output.".format(signature_name))
return
if len(candidates) > 1:
print("Found multiple equivalent outputs '{}'. Unable to rename.".format(",".join([ o.name for o in candidates ])))
return
graph_output = candidates[0]
print("Found equivalent output '{}'. Assuming this is correct.".format(graph_output.name))
if graph_output.name == signature_name:
print("Signature output '{}' already exists. Skipping.".format(signature_name))
return
output_node = find(onnx_model.graph.node, lambda node: graph_output.name in node.output)
if output_node is None:
print("Node generating graph output '{}' was not found. Unable to rename.".format(graph_output.name))
return
print("Renamed output from '{}' to '{}'".format(graph_output.name, signature_name))
output_node.output[index_of(output_node.output, graph_output.name)] = signature_name
graph_output.name = signature_name
def verify_outputs(args, onnx_model):
tag_sets = saved_model_utils.get_saved_model_tag_sets(args.saved_model)
for tag_set in tag_sets:
tag_set = ','.join(tag_set)
meta_graph_def = saved_model_utils.get_meta_graph_def(args.saved_model, tag_set)
signature_def_map = meta_graph_def.signature_def
for signature_def_key in signature_def_map.keys():
outputs_tensor_info = signature_def_map[signature_def_key].outputs
for output_key, output_tensor in outputs_tensor_info.items():
rename_output(onnx_model, output_key, output_tensor)
print("Inputs in model: {}".format(", ".join(["'{}'".format(o.name) for o in onnx_model.graph.input if not has_initializer(onnx_model, o.name)])))
print("Outputs in model: {}".format(", ".join(["'{}'".format(o.name) for o in onnx_model.graph.output])))
def main():
args = convert.get_args()
convert.main()
onnx_model = onnx.load(args.output)
verify_outputs(args, onnx_model)
onnx.save(onnx_model, args.output)
if __name__ == "__main__":
main()
|
