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

Guo, Yejun committed 7 years ago

Commit 2c014341 parent aac382e2 files · +12 −7Browse files

Unified · Split

2 files

tools/python/convert.py +5 −1
tools/python/convert_from_tensorflow.py +7 −6

@@ -27,6 +27,7 @@ def get_arguments(): · 26 unmodified lines
27	27	parser.add_argument('--outdir', type=str, default='./', help='where to put generated files')
28	28	parser.add_argument('--infmt', type=str, default='tensorflow', help='format of the deep learning model')
29	29	parser.add_argument('infile', help='path to the deep learning model with weights')
	30	+ parser.add_argument('--dump4tb', type=str, default='no', help='dump file for visualization in tensorboard')
30	31
31	32	return parser.parse_args()
32	33
@@ -44,9 +45,12 @@ def main(): · 11 unmodified lines
44	45	basefile = os.path.split(args.infile)[1]
45	46	basefile = os.path.splitext(basefile)[0]
46	47	outfile = os.path.join(args.outdir, basefile) + '.model'
	48	+ dump4tb = False
	49	+ if args.dump4tb.lower() in ('yes', 'true', 't', 'y', '1'):
	50	+ dump4tb = True
47	51
48	52	if args.infmt == 'tensorflow':
49		− convert_from_tensorflow(args.infile, outfile)
	53	+ convert_from_tensorflow(args.infile, outfile, dump4tb)
50	54
51	55	if __name__ == '__main__':
52	56	main()

tools/python/convert_from_tensorflow.py+7−6

@@ -24,10 +24,11 @@ import sys, struct · 23 unmodified lines
24	24	__all__ = ['convert_from_tensorflow']
25	25
26	26	class TFConverter:
27		− def __init__(self, graph_def, nodes, outfile):
	27	+ def __init__(self, graph_def, nodes, outfile, dump4tb):
28	28	self.graph_def = graph_def
29	29	self.nodes = nodes
30	30	self.outfile = outfile
	31	+ self.dump4tb = dump4tb
31	32	self.layer_number = 0
32	33	self.output_names = []
33	34	self.name_node_dict = {}
@@ -42,8 +43,8 @@ class TFConverter: · 8 unmodified lines
42	43	def dump_for_tensorboard(self):
43	44	graph = tf.get_default_graph()
44	45	tf.import_graph_def(self.graph_def, name="")
45		− # tensorboard --logdir=/tmp/graph
46	46	tf.summary.FileWriter('/tmp/graph', graph)
	47	+ print('graph saved, run "tensorboard --logdir=/tmp/graph" to see it')
47	48
48	49
49	50	def get_conv2d_params(self, node):
@@ -197,18 +198,18 @@ class TFConverter: · 147 unmodified lines
197	198	self.remove_identity()
198	199	self.generate_edges()
199	200
200		− #check the graph with tensorboard with human eyes
201		− #self.dump_for_tensorboard()
	201	+ if self.dump4tb:
	202	+ self.dump_for_tensorboard()
202	203
203	204	self.dump_to_file()
204	205
205	206
206		−def convert_from_tensorflow(infile, outfile):
	207	+def convert_from_tensorflow(infile, outfile, dump4tb):
207	208	with open(infile, 'rb') as f:
208	209	# read the file in .proto format
209	210	graph_def = tf.GraphDef()
210	211	graph_def.ParseFromString(f.read())
211	212	nodes = graph_def.node
212	213
213		− converter = TFConverter(graph_def, nodes, outfile)
	214	+ converter = TFConverter(graph_def, nodes, outfile, dump4tb)
214	215	converter.run()

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

Keyboard shortcuts

General

Issues and pull requests

Issues, pull requests and Inbox

Repository code