Skip to content

Commit 0390798

Browse files
nbcsmnbcsm
committed
unify debug mode
1 parent f0d5b4b commit 0390798

13 files changed

+74
-38
lines changed

‎tests/common.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@ def __init__(self):
3131
self.backend = os.environ.get("TF2ONNX_TEST_BACKEND", "onnxruntime")
3232
self.backend_version = self._get_backend_version()
3333
self.log_level = logging.WARNING
34-
self.is_debug_mode = False
3534
self.temp_dir = utils.get_temp_directory()
3635

3736
@property
@@ -46,6 +45,10 @@ def is_onnxruntime_backend(self):
4645
def is_caffe2_backend(self):
4746
return self.backend == "caffe2"
4847

48+
@property
49+
def is_debug_mode(self):
50+
return utils.is_debug_mode()
51+
4952
def _get_tf_version(self):
5053
import tensorflow as tf
5154
return LooseVersion(tf.__version__)
@@ -92,11 +95,13 @@ def load():
9295
parser.add_argument("unittest_args", nargs='*')
9396

9497
args = parser.parse_args()
98+
if args.debug:
99+
utils.set_debug_mode(True)
100+
95101
config.backend = args.backend
96102
config.opset = args.opset
97103
config.target = args.target.split(',')
98104
config.log_level = logging.get_verbosity_level(args.verbose, config.log_level)
99-
config.is_debug_mode = args.debug
100105
if args.temp_dir:
101106
config.temp_dir = args.temp_dir
102107

‎tests/run_pretrained_models.py

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ def create_onnx_file(name, model_proto, inputs, outdir):
203203
utils.save_protobuf(model_path, model_proto)
204204
print("\tcreated", model_path)
205205

206-
def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=None, extra_opset=None,
206+
def run_test(self, name, backend="caffe2", onnx_file=None, opset=None, extra_opset=None,
207207
perf=None, fold_const=None):
208208
"""Run complete test against backend."""
209209
print(name)
@@ -241,7 +241,7 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
241241
inputs[k] = v
242242

243243
graph_def = tf2onnx.tfonnx.tf_optimize(inputs.keys(), self.output_names, graph_def, fold_const)
244-
if debug:
244+
if utils.is_debug_mode():
245245
utils.save_protobuf(os.path.join(TEMP_DIR, name + "_after_tf_optimize.pb"), graph_def)
246246
shape_override = {}
247247
g = tf.import_graph_def(graph_def, name='')
@@ -271,13 +271,13 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
271271
onnx_graph = self.to_onnx(sess.graph, opset=opset, extra_opset=extra_opset,
272272
shape_override=shape_override, input_names=inputs.keys())
273273
model_proto = onnx_graph.make_model("converted from tf2onnx")
274-
new_model_proto = optimizer.optimize_graph(onnx_graph, debug=debug).make_model("optimized")
274+
new_model_proto = optimizer.optimize_graph(onnx_graph).make_model("optimized")
275275
if new_model_proto:
276276
model_proto = new_model_proto
277277
else:
278278
print("\tNON-CRITICAL, optimizers are not applied successfully")
279279
print("\tto_onnx", "OK")
280-
if debug:
280+
if utils.is_debug_mode():
281281
onnx_graph.dump_graph()
282282
if onnx_file:
283283
self.create_onnx_file(name, model_proto, inputs, onnx_file)
@@ -333,7 +333,7 @@ def get_args():
333333
parser.add_argument("--extra_opset", default=None,
334334
help="extra opset with format like domain:version, e.g. com.microsoft:1")
335335
parser.add_argument("--verbose", "-v", help="verbose output, option is additive", action="count")
336-
parser.add_argument("--debug", help="debug vlog", action="store_true")
336+
parser.add_argument("--debug", help="debug mode", action="store_true")
337337
parser.add_argument("--list", help="list tests", action="store_true")
338338
parser.add_argument("--onnx-file", help="create onnx file in directory")
339339
parser.add_argument("--perf", help="capture performance numbers")
@@ -372,6 +372,8 @@ def tests_from_yaml(fname):
372372
def main():
373373
args = get_args()
374374
logging.basicConfig(level=logging.get_verbosity_level(args.verbose))
375+
if args.debug:
376+
utils.set_debug_mode(True)
375377

376378
Test.cache_dir = args.cache
377379
Test.target = args.target
@@ -392,15 +394,15 @@ def main():
392394
continue
393395
count += 1
394396
try:
395-
ret = t.run_test(test, backend=args.backend, debug=args.debug, onnx_file=args.onnx_file,
397+
ret = t.run_test(test, backend=args.backend, onnx_file=args.onnx_file,
396398
opset=args.opset, extra_opset=args.extra_opset, perf=args.perf,
397399
fold_const=args.fold_const)
398400
except Exception as ex:
399401
ret = None
400402
tb = traceback.format_exc()
401403
print(ex, tb)
402404
finally:
403-
if not args.debug:
405+
if not utils.is_debug_mode():
404406
utils.delete_directory(TEMP_DIR)
405407
if not ret:
406408
failed += 1

‎tests/test_optimizers.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from backend_test_base import Tf2OnnxBackendTestBase
1515
from common import unittest_main, group_nodes_by_type
1616

17+
1718
# pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
1819

1920
class OptimizerTests(Tf2OnnxBackendTestBase):
@@ -34,8 +35,8 @@ def run_and_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
3435
current = GraphUtil.get_node_count_from_onnx_graph(new_proto.graph)
3536

3637
self.assertTrue(current[op_type] == remaining_op_num,
37-
msg="Expect " + str(remaining_op_num) + " " + op_type + " ops left, but actually " +
38-
str(current[op_type]) + " left")
38+
msg="Expect " + str(remaining_op_num) + " " + op_type + " ops left, but actually " + str(
39+
current[op_type]) + " left")
3940

4041
if self.config.is_onnxruntime_backend:
4142
expected = self.run_onnxruntime(origin_model_path, onnx_feed_dict, output_names_with_port)
@@ -179,14 +180,13 @@ def test_trans_output_as_graph_outputs(self):
179180
graph.outputs = [identity_op.input[0]]
180181
graph.remove_node(identity_op.name)
181182

182-
optimized_graph = GraphUtil.optimize_graph(graph, "onnx-tests")
183+
optimized_graph = GraphUtil.optimize_graph(graph)
183184

184185
self.assertTrue(optimized_graph, msg="graph after optimizer should not be None")
185186

186187
trans_cnt = len(group_nodes_by_type(optimized_graph)["Transpose"])
187188

188-
self.assertTrue(trans_cnt == 1, msg="Expect 1 Transpose ops left, but actually " +
189-
str(trans_cnt) + " left")
189+
self.assertTrue(trans_cnt == 1, msg="Expect 1 Transpose ops left, but actually " + str(trans_cnt) + " left")
190190

191191
# Tranpose Optimizer Tests End
192192

@@ -292,7 +292,7 @@ def test_identity_in_subgraph_non_graph_output(self):
292292
],
293293
[helper.make_tensor_value_info("loop_cond_output", TensorProto.BOOL, ()),
294294
helper.make_tensor_value_info("loop_var_out_1", TensorProto.FLOAT, ())
295-
],
295+
],
296296
)
297297
# sub graph ends
298298

@@ -394,7 +394,7 @@ def test_duplicated_different_output_length(self):
394394

395395
model_proto = helper.make_model(graph, producer_name="onnx-tests")
396396
self.run_merge_duplicated_nodes_compare(["value1", "mask", "value2"],
397-
{"X": np.random.randn(5,).astype(np.float32)},
397+
{"X": np.random.randn(5).astype(np.float32)},
398398
model_proto,
399399
op_type="Dropout", remaining_op_num=2)
400400

@@ -417,9 +417,10 @@ def test_duplicated_need_multiple_run(self):
417417
)
418418

419419
model_proto = helper.make_model(graph, producer_name="onnx-tests")
420-
self.run_merge_duplicated_nodes_compare(["res"], {"X": np.random.randn(5,).astype(np.float32)},
420+
self.run_merge_duplicated_nodes_compare(["res"], {"X": np.random.randn(5).astype(np.float32)},
421421
model_proto,
422422
op_type="Log", remaining_op_num=3)
423+
423424
# Merge Duplicated Nodes Optimizer Tests End
424425

425426
# Const Fold Optimizer Tests Start

‎tf2onnx/constants.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,3 +34,6 @@
3434
NHWC_TO_NCHW = [0, 3, 1, 2]
3535
HWCN_TO_NCHW = [3, 2, 0, 1]
3636
NCHW_TO_HWCN = [2, 3, 1, 0]
37+
38+
# Environment variables
39+
ENV_TF2ONNX_DEBUG_MODE = "TF2ONNX_DEBUG_MODE"

‎tf2onnx/convert.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ def get_args():
3838
help="target platform")
3939
parser.add_argument("--continue_on_error", help="continue_on_error", action="store_true")
4040
parser.add_argument("--verbose", "-v", help="verbose output, option is additive", action="count")
41+
parser.add_argument("--debug", help="debug mode", action="store_true")
4142
parser.add_argument("--fold_const", help="enable tf constant_folding transformation before conversion",
4243
action="store_true")
4344
# experimental
@@ -81,6 +82,8 @@ def default_custom_op_handler(ctx, node, name, args):
8182
def main():
8283
args = get_args()
8384
logging.basicConfig(level=logging.get_verbosity_level(args.verbose))
85+
if args.debug:
86+
utils.set_debug_mode(True)
8487

8588
# override unknown dimensions from -1 to 1 (aka batchsize 1) since not every runtime does
8689
# support unknown dimensions.

‎tf2onnx/graph.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1057,11 +1057,11 @@ class GraphUtil(object):
10571057
"""Utilities for Graph manipulation."""
10581058

10591059
@staticmethod
1060-
def optimize_graph(graph, debug=False):
1061-
return optimizer.optimize_graph(graph, debug)
1060+
def optimize_graph(graph):
1061+
return optimizer.optimize_graph(graph)
10621062

10631063
@staticmethod
1064-
def optimize_model_proto(onnx_model_proto, debug=False):
1064+
def optimize_model_proto(onnx_model_proto):
10651065
"""Optimize the model proto, for example: eliminating all useless Transpose pairs.
10661066
10671067
Returns:
@@ -1071,7 +1071,7 @@ def optimize_model_proto(onnx_model_proto, debug=False):
10711071
try:
10721072
kwargs = GraphUtil.get_onnx_model_properties(onnx_model_proto)
10731073
graph = GraphUtil.create_graph_from_onnx_model(onnx_model_proto)
1074-
graph = GraphUtil.optimize_graph(graph, debug)
1074+
graph = GraphUtil.optimize_graph(graph)
10751075
model_proto = graph.make_model(onnx_model_proto.graph.doc_string,
10761076
graph_name=onnx_model_proto.graph.name, **kwargs)
10771077

‎tf2onnx/optimizer/__init__.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,11 @@
2828
])
2929

3030

31-
def optimize_graph(graph, debug=False):
31+
def optimize_graph(graph):
3232
try:
3333
opts = _get_optimizers()
3434
for opt in opts.values():
35-
graph = opt(debug=debug).optimize(graph)
35+
graph = opt().optimize(graph)
3636

3737
graph.update_proto()
3838
return graph

‎tf2onnx/optimizer/const_fold_optimizer.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@ def _internal_fun(func):
2626

2727
class ConstFoldOptimizer(GraphOptimizerBase):
2828

29-
def __init__(self, debug=False): # pylint: disable=useless-super-delegation
30-
super(ConstFoldOptimizer, self).__init__(debug)
29+
def __init__(self): # pylint: disable=useless-super-delegation
30+
super(ConstFoldOptimizer, self).__init__()
3131

3232
def _optimize(self, graph):
3333
return self._apply_optimization(graph, self._optimize_at_current_graph_level)

‎tf2onnx/optimizer/identity_optimizer.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616
class IdentityOptimizer(GraphOptimizerBase):
1717
"""Identity Optimizer."""
1818

19-
def __init__(self, debug=False):
20-
super(IdentityOptimizer, self).__init__(debug)
19+
def __init__(self):
20+
super(IdentityOptimizer, self).__init__()
2121
self._g = None
2222

2323
def optimize(self, graph):

‎tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@ class MergeDuplicatedNodesOptimizer(GraphOptimizerBase):
2020
"""Remove duplicate nodes.
2121
"""
2222

23-
def __init__(self, debug=False):
24-
super(MergeDuplicatedNodesOptimizer, self).__init__(debug)
23+
def __init__(self):
24+
super(MergeDuplicatedNodesOptimizer, self).__init__()
2525
# used internally
2626
self._graph_can_be_optimized = True
2727

0 commit comments

Comments
 (0)