From 2c01434d60dbc607076a34d129e93a3585cbb277 Mon Sep 17 00:00:00 2001 From: "Guo, Yejun" Date: Tue, 30 Jul 2019 09:25:56 +0800 Subject: [PATCH] convert_from_tensorflow.py: add option to dump graph for visualization in tensorboard Signed-off-by: Guo, Yejun Signed-off-by: Pedro Arthur --- tools/python/convert.py | 6 +++++- tools/python/convert_from_tensorflow.py | 13 +++++++------ 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/tools/python/convert.py b/tools/python/convert.py index 662b429066..64cf76b2d8 100644 --- a/tools/python/convert.py +++ b/tools/python/convert.py @@ -27,6 +27,7 @@ def get_arguments(): parser.add_argument('--outdir', type=str, default='./', help='where to put generated files') parser.add_argument('--infmt', type=str, default='tensorflow', help='format of the deep learning model') parser.add_argument('infile', help='path to the deep learning model with weights') + parser.add_argument('--dump4tb', type=str, default='no', help='dump file for visualization in tensorboard') return parser.parse_args() @@ -44,9 +45,12 @@ def main(): basefile = os.path.split(args.infile)[1] basefile = os.path.splitext(basefile)[0] outfile = os.path.join(args.outdir, basefile) + '.model' + dump4tb = False + if args.dump4tb.lower() in ('yes', 'true', 't', 'y', '1'): + dump4tb = True if args.infmt == 'tensorflow': - convert_from_tensorflow(args.infile, outfile) + convert_from_tensorflow(args.infile, outfile, dump4tb) if __name__ == '__main__': main() diff --git a/tools/python/convert_from_tensorflow.py b/tools/python/convert_from_tensorflow.py index 041c82cf34..804c14f5f2 100644 --- a/tools/python/convert_from_tensorflow.py +++ b/tools/python/convert_from_tensorflow.py @@ -24,10 +24,11 @@ import sys, struct __all__ = ['convert_from_tensorflow'] class TFConverter: - def __init__(self, graph_def, nodes, outfile): + def __init__(self, graph_def, nodes, outfile, dump4tb): self.graph_def = graph_def self.nodes = nodes self.outfile = outfile + self.dump4tb = dump4tb self.layer_number = 0 self.output_names = [] self.name_node_dict = {} @@ -42,8 +43,8 @@ class TFConverter: def dump_for_tensorboard(self): graph = tf.get_default_graph() tf.import_graph_def(self.graph_def, name="") - # tensorboard --logdir=/tmp/graph tf.summary.FileWriter('/tmp/graph', graph) + print('graph saved, run "tensorboard --logdir=/tmp/graph" to see it') def get_conv2d_params(self, node): @@ -197,18 +198,18 @@ class TFConverter: self.remove_identity() self.generate_edges() - #check the graph with tensorboard with human eyes - #self.dump_for_tensorboard() + if self.dump4tb: + self.dump_for_tensorboard() self.dump_to_file() -def convert_from_tensorflow(infile, outfile): +def convert_from_tensorflow(infile, outfile, dump4tb): with open(infile, 'rb') as f: # read the file in .proto format graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) nodes = graph_def.node - converter = TFConverter(graph_def, nodes, outfile) + converter = TFConverter(graph_def, nodes, outfile, dump4tb) converter.run()