Signed-off-by: Guo, Yejun <yejun....@intel.com>
---
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 662b429..64cf76b 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 041c82c..804c14f 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()
--
2.7.4
_______________________________________________
ffmpeg-devel mailing list
ffmpeg-devel@ffmpeg.org
https://ffmpeg.org/mailman/listinfo/ffmpeg-devel
To unsubscribe, visit link above, or email
ffmpeg-devel-requ...@ffmpeg.org with subject "unsubscribe".
