forked from kerlomz/captcha_trainer
-
Notifications
You must be signed in to change notification settings - Fork 0
/
tf_onnx_util.py
132 lines (102 loc) · 4.78 KB
/
tf_onnx_util.py
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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Author: kerlomz <kerlomz@gmail.com>
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT license.
"""
python -m tf2onnx.convert : tool to convert a frozen tensorflow graph to onnx
"""
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import argparse
import sys
import tensorflow as tf
from tf2onnx.tfonnx import process_tf_graph, tf_optimize
from tf2onnx import constants, logging, utils, optimizer
from tf_graph_util import convert_variables_to_constants
# from tensorflow.python.framework.graph_util import convert_variables_to_constants
# pylint: disable=unused-argument
_HELP_TEXT = """
Usage Examples:
python -m tf2onnx.convert --saved-model saved_model_dir --output model.onnx
python -m tf2onnx.convert --input frozen_graph.pb --inputs X:0 --outputs output:0 --output model.onnx
python -m tf2onnx.convert --checkpoint checkpoint.meta --inputs X:0 --outputs output:0 --output model.onnx
For help and additional information see:
https://github.com/onnx/tensorflow-onnx
If you run into issues, open an issue here:
https://github.com/onnx/tensorflow-onnx/issues
"""
logger = tf.compat.v1.logging
# logger = logging.getLogger(constants.TF2ONNX_PACKAGE_NAME)
def freeze_session(sess, keep_var_names=None, output_names=None, clear_devices=True):
"""Freezes the state of a session into a pruned computation graph."""
output_names = [i.split(':')[:-1][0] for i in output_names]
graph = sess.graph
with graph.as_default():
freeze_var_names = list(set(v.op.name for v in tf.compat.v1.global_variables()).difference(keep_var_names or []))
output_names = output_names or []
output_names += [v.op.name for v in tf.compat.v1.global_variables()]
input_graph_def = graph.as_graph_def(add_shapes=True)
if clear_devices:
for node in input_graph_def.node:
node.device = ""
frozen_graph = convert_variables_to_constants(sess, input_graph_def, output_names, freeze_var_names)
return frozen_graph
def remove_redundant_inputs(frozen_graph, input_names):
"""Remove redundant inputs not in frozen graph."""
frozen_inputs = []
# get inputs in frozen graph
for n in frozen_graph.node:
for inp in input_names:
if utils.node_name(inp) == n.name:
frozen_inputs.append(inp)
deleted_inputs = list(set(input_names) - set(frozen_inputs))
if deleted_inputs:
logger.warning("inputs [%s] is not in frozen graph, delete them", ",".join(deleted_inputs))
return frozen_inputs
def from_graphdef(sess, graph_def, model_path, input_names, output_names):
"""Load tensorflow graph from graphdef."""
# make sure we start with clean default graph
with tf.io.gfile.GFile(model_path, 'rb') as f:
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
frozen_graph = freeze_session(sess, output_names=output_names)
input_names = remove_redundant_inputs(frozen_graph, input_names)
# clean up
return frozen_graph, input_names, output_names
def convert_onnx(sess, graph_def, input_path, inputs_op, outputs_op):
graphdef = input_path
if inputs_op:
inputs_op, shape_override = utils.split_nodename_and_shape(inputs_op)
if outputs_op:
outputs_op = outputs_op.split(",")
# logging.basicConfig(level=logging.get_verbosity_level(True))
utils.set_debug_mode(True)
graph_def, inputs_op, outputs_op = from_graphdef(sess, graph_def, graphdef, inputs_op, outputs_op)
model_path = graphdef
graph_def = tf_optimize(inputs_op, outputs_op, graph_def, True)
with tf.Graph().as_default() as tf_graph:
tf.compat.v1.import_graph_def(graph_def, name='')
with tf.compat.v1.Session(graph=tf_graph):
g = process_tf_graph(tf_graph,
continue_on_error=True,
target=",".join(constants.DEFAULT_TARGET),
opset=11,
custom_op_handlers=None,
extra_opset=None,
shape_override=None,
input_names=inputs_op,
output_names=outputs_op,
inputs_as_nchw=None)
onnx_graph = optimizer.optimize_graph(g)
model_proto = onnx_graph.make_model("converted from {}".format(model_path))
# write onnx graph
logger.info("")
logger.info("Successfully converted TensorFlow model %s to ONNX", model_path)
# if args.output:
output_path = input_path.replace(".pb", ".onnx")
utils.save_protobuf(output_path, model_proto)
logger.info("ONNX model is saved at %s", output_path)
if __name__ == "__main__":
pass