From 799935d5b39e45cf0f00a5de16d61fdd328552be Mon Sep 17 00:00:00 2001 From: =?utf8?q?=EC=9C=A4=ED=98=84=EC=8B=9D/=EB=8F=99=EC=9E=91=EC=A0=9C?= =?utf8?q?=EC=96=B4Lab=28SR=29/Principal=20Engineer/=EC=82=BC=EC=84=B1?= =?utf8?q?=EC=A0=84=EC=9E=90?= Date: Wed, 11 Jul 2018 15:52:18 +0900 Subject: [PATCH] Prints summary of pb file (#1946) Added `--summary` option to prints the list of operations and their counts in pb file. Written for #1927. Signed-off-by: Hyun Sik Yoon --- tools/pbfile_tool/pb_info.py | 36 +++++++++++++++++++++++++++++++++++- tools/pbfile_tool/readme.md | 4 +++- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/tools/pbfile_tool/pb_info.py b/tools/pbfile_tool/pb_info.py index cda34ec..8101b05 100755 --- a/tools/pbfile_tool/pb_info.py +++ b/tools/pbfile_tool/pb_info.py @@ -95,6 +95,35 @@ def print_graph_info(pb_path, optype_substring): op_count += 1 +def print_summary(pb_path, optype_substring): + op_map = {} + with tf.Session() as sess: + importGraphIntoSession(sess, pb_path) + + op_count = 1 + graph = sess.graph + ops = graph.get_operations() + for op in ops: + process = False + if optype_substring == "*": + process = True + elif op.type.lower().find(optype_substring.lower()) != -1: + process = True + + if process: + if op_map.get(op.type) == None: + op_map[op.type] = 1 + else: + op_map[op.type] += 1 + + # print op list + print("") + print("Total number of operation types : " + str(len(op_map.keys()))) + print("") + for op_type, count in op_map.items(): + print("\t" + op_type + " : \t" + str(count)) + + if __name__ == "__main__": parser = argparse.ArgumentParser(description='Prints information inside pb file') @@ -103,7 +132,12 @@ if __name__ == "__main__": parser.add_argument( "op_subst", help="substring of operations. only info of these operasions will be printed.") + parser.add_argument( + "--summary", help="print summary of operations", action="store_true") args = parser.parse_args() - print_graph_info(args.pb_file, args.op_subst) + if args.summary: + print_summary(args.pb_file, args.op_subst) + else: + print_graph_info(args.pb_file, args.op_subst) diff --git a/tools/pbfile_tool/readme.md b/tools/pbfile_tool/readme.md index 30e0fa4..15594e6 100644 --- a/tools/pbfile_tool/readme.md +++ b/tools/pbfile_tool/readme.md @@ -5,4 +5,6 @@ - first arg: pb file - second arg: substring of operation. Only operations that has "conv" substring as its type will be printed. (case-insensitive) - `./tools/pbfile_tool/pb_info.py pbfile_path "*"` - - pass "*" as the second param to print all operations + - pass "*" as the second param to print all operations + - `./tools/pbfile_tool/pb_info.py pbfile_path "*" --summary` + - prints the list of operations and their counts -- 2.7.4