Showing the changes to tools/python/convert_from_tensorflow.py (1 of 2 files). Show all
Unified · Split
| @@ -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() |