ffmpeg | branch: master | Guo, Yejun <yejun....@intel.com> | Tue Jul 30 
09:25:56 2019 +0800| [2c01434d60dbc607076a34d129e93a3585cbb277] | committer: 
Pedro Arthur

convert_from_tensorflow.py: add option to dump graph for visualization in 
tensorboard

Signed-off-by: Guo, Yejun <yejun....@intel.com>
Signed-off-by: Pedro Arthur <bygran...@gmail.com>

> http://git.videolan.org/gitweb.cgi/ffmpeg.git/?a=commit;h=2c01434d60dbc607076a34d129e93a3585cbb277
---

 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()

_______________________________________________
ffmpeg-cvslog mailing list
ffmpeg-cvslog@ffmpeg.org
https://ffmpeg.org/mailman/listinfo/ffmpeg-cvslog

To unsubscribe, visit link above, or email
ffmpeg-cvslog-requ...@ffmpeg.org with subject "unsubscribe".

Reply via email to