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