virtualenv/FFmpeg
1
0
Fork 0
mirror of https://github.com/FFmpeg/FFmpeg.git synced 2024-12-23 12:43:46 +02:00
Code Issues Releases Activity

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

Browse Source 
Signed-off-by: Guo, Yejun <yejun.guo@intel.com>
Signed-off-by: Pedro Arthur <bygrandao@gmail.com>
This commit is contained in:
Guo, Yejun 2019-07-30 09:25:56 +08:00 committed by Pedro Arthur
parent aac382e9e5
commit 2c01434d60
2 changed files with 12 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tools/python/convert.py
View File

@ -27,6 +27,7 @@ def get_arguments():
parser.add_argument('--outdir', type=str, default='./', help='where to put generated files') 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('--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('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() return parser.parse_args()
@ -44,9 +45,12 @@ def main():
basefile = os.path.split(args.infile)[1] basefile = os.path.split(args.infile)[1]
basefile = os.path.splitext(basefile)[0] basefile = os.path.splitext(basefile)[0]
outfile = os.path.join(args.outdir, basefile) + '.model' 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': if args.infmt == 'tensorflow':
convert_from_tensorflow(args.infile, outfile) convert_from_tensorflow(args.infile, outfile, dump4tb)
if __name__ == '__main__': if __name__ == '__main__':
main() main()

13
tools/python/convert_from_tensorflow.py
View File

@ -24,10 +24,11 @@ import sys, struct
__all__ = ['convert_from_tensorflow'] __all__ = ['convert_from_tensorflow']
class TFConverter: class TFConverter:
def __init__(self, graph_def, nodes, outfile): def __init__(self, graph_def, nodes, outfile, dump4tb):
self.graph_def = graph_def self.graph_def = graph_def
self.nodes = nodes self.nodes = nodes
self.outfile = outfile self.outfile = outfile
self.dump4tb = dump4tb
self.layer_number = 0 self.layer_number = 0
self.output_names = [] self.output_names = []
self.name_node_dict = {} self.name_node_dict = {}
@ -42,8 +43,8 @@ class TFConverter:
def dump_for_tensorboard(self): def dump_for_tensorboard(self):
graph = tf.get_default_graph() graph = tf.get_default_graph()
tf.import_graph_def(self.graph_def, name="") tf.import_graph_def(self.graph_def, name="")
# tensorboard --logdir=/tmp/graph
tf.summary.FileWriter('/tmp/graph', graph) tf.summary.FileWriter('/tmp/graph', graph)
print('graph saved, run "tensorboard --logdir=/tmp/graph" to see it')
def get_conv2d_params(self, node): def get_conv2d_params(self, node):
@ -197,18 +198,18 @@ class TFConverter:
self.remove_identity() self.remove_identity()
self.generate_edges() self.generate_edges()
#check the graph with tensorboard with human eyes if self.dump4tb:
#self.dump_for_tensorboard() self.dump_for_tensorboard()
self.dump_to_file() self.dump_to_file()
def convert_from_tensorflow(infile, outfile): def convert_from_tensorflow(infile, outfile, dump4tb):
with open(infile, 'rb') as f: with open(infile, 'rb') as f:
# read the file in .proto format # read the file in .proto format
graph_def = tf.GraphDef() graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read()) graph_def.ParseFromString(f.read())
nodes = graph_def.node nodes = graph_def.node
converter = TFConverter(graph_def, nodes, outfile) converter = TFConverter(graph_def, nodes, outfile, dump4tb)
converter.run() converter.run()