diff options
| author | Guo, Yejun <yejun.guo@intel.com> | 2019-07-30 09:25:56 +0800 |
|---|---|---|
| committer | Pedro Arthur <bygrandao@gmail.com> | 2019-08-15 14:58:19 -0300 |
| commit | 2c01434d60dbc607076a34d129e93a3585cbb277 (patch) | |
| tree | cb409a1563eefe4a20405298088ab361aaccb44f /tools/python/convert_from_tensorflow.py | |
| parent | aac382e9e5f128898a19cc98259a25d29d93df71 (diff) | |
| download | ffmpeg-2c01434d60dbc607076a34d129e93a3585cbb277.tar.gz | |
convert_from_tensorflow.py: add option to dump graph for visualization in tensorboard
Signed-off-by: Guo, Yejun <yejun.guo@intel.com>
Signed-off-by: Pedro Arthur <bygrandao@gmail.com>
Diffstat (limited to 'tools/python/convert_from_tensorflow.py')
| -rw-r--r-- | tools/python/convert_from_tensorflow.py | 13 |
1 files changed, 7 insertions, 6 deletions
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() |